Caffe2 - C++ API
A deep learning, cross platform ML framework
shape_op.cc
1 #include "caffe2/operators/shape_op.h"
2 
3 namespace caffe2 {
4 
5 REGISTER_CPU_OPERATOR(Shape, ShapeOp<CPUContext>);
6 
7 OPERATOR_SCHEMA(Shape)
8  .NumInputs(1)
9  .NumOutputs(1)
10  .Arg(
11  "axes",
12  "*(type: int[])* Array of interested axes."
13  "If given, this operator only returns the dimensions of the given axes."
14  "Otherwise, the operator returns the dimensions of all axes.")
15  .TensorInferenceFunction([](const OperatorDef& def,
16  const vector<TensorShape>& in) {
17  ArgumentHelper args(def);
18  const vector<int>& axes = args.GetRepeatedArgument<int>("axes");
19  vector<TensorShape> out(1);
20  if (axes.empty()) {
21  out[0].add_dims(in[0].dims().size());
22  } else {
23  out[0].add_dims(axes.size());
24  }
25  out[0].set_data_type(TensorProto::INT64);
26  return out;
27  })
28  .SetDoc(R"DOC(
29 Produce a 1D int64 tensor with the shape of the input tensor.
30 If called with an optional argument `axes`, the result will only
31 contain the dimensions of specified axes.
32 
33 Github Link:
34 - https://github.com/pytorch/pytorch/blob/master/caffe2/operators/shape_op.cc
35 
36 <details>
37 
38 <summary> <b>Example</b> </summary>
39 
40 **Code**
41 
42 ```
43 
44 workspace.ResetWorkspace()
45 
46 op = core.CreateOperator(
47  "Shape",
48  ["X"],
49  ["shape"],
50 )
51 
52 workspace.FeedBlob("X", (np.random.randint(10, size=(2,3))))
53 print("X:", workspace.FetchBlob("X"))
54 workspace.RunOperatorOnce(op)
55 print("shape:", workspace.FetchBlob("shape"))
56 
57 ```
58 
59 **Result**
60 
61 ```
62 
63 X:
64 [[3 2 5]
65  [5 7 3]]
66 shape: [2 3]
67 
68 ```
69 
70 </details>
71 
72  )DOC")
73  .Input(0,"X", "*(type: Tensor)* Input tensor.")
74  .Output(0,"shape", "*(type: Tensor)* Output tensor containing shape of input tensor.");
75 
76 SHOULD_NOT_DO_GRADIENT(Shape);
77 
78 } // namespace caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13