Caffe2 - C++ API
A deep learning, cross platform ML framework
shape_op.cc
1 
17 #include "caffe2/operators/shape_op.h"
18 
19 namespace caffe2 {
20 
21 REGISTER_CPU_OPERATOR(Shape, ShapeOp<CPUContext>);
22 
23 OPERATOR_SCHEMA(Shape)
24  .NumInputs(1)
25  .NumOutputs(1)
26  .TensorInferenceFunction([](const OperatorDef& /*def*/,
27  const vector<TensorShape>& in) {
28  vector<TensorShape> out(1);
29  out[0].add_dims(in[0].dims().size());
30  out[0].set_data_type(TensorProto::INT32);
31  return out;
32  })
33  .SetDoc("Produce a 1D int64 tensor with the shape of the input tensor.");
34 
35 SHOULD_NOT_DO_GRADIENT(Shape);
36 
37 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.