1 #include "caffe2/operators/reshape_op.h" 2 #include "caffe2/utils/math.h" 6 REGISTER_CPU_OPERATOR(Reshape, ReshapeOp<float, CPUContext>);
8 OPERATOR_SCHEMA(Reshape)
11 .TensorInferenceFunction(
12 [](
const OperatorDef& def,
const vector<TensorShape>& in) {
13 vector<TensorShape> out(2);
16 out[1].set_data_type(TensorProto::INT64);
17 out[1].add_dims(in[0].dims_size());
19 ArgumentHelper helper(def);
20 if (!helper.HasArgument(
"shape")) {
25 "New shape must be specified by either the input blob or the " 27 out[0].set_unknown_shape(
true);
33 "New shape must not be specified by the input blob and the " 34 "argument `shape` at the same time.");
37 auto actualNewShape = helper.GetRepeatedArgument<int64_t>(
"shape");
41 for (
int i = 0; i < actualNewShape.size(); ++i) {
45 "The dimensions in argument `shape` " 46 "must not be a negative number.");
48 if (actualNewShape[i] == 0) {
52 "Argument `shape` has a dimension set to zero that exceeds " 53 "the original dimension size.");
54 actualNewShape[i] = in[0].dims(i);
60 int64_t totalSize = 1;
61 for (
const auto d : in[0].dims()) {
66 for (
int i = 0; i < actualNewShape.size(); ++i) {
67 const auto dim = actualNewShape[i];
71 "Argument `shape` has more than one missing dimension.");
78 if (unknownIdx != -1) {
80 totalSize % size == 0,
81 "Argument `shape` does not agree with the input data.",
87 actualNewShape[unknownIdx] = totalSize / size;
92 "Argument `shape` does not agree with the input data.",
100 out[0].set_data_type(in[0].data_type());
101 for (
const auto d : actualNewShape) {
106 .AllowInplace({{0, 0}})
108 Reshape the input tensor similar to numpy.reshape. 110 It takes a tensor as input and an optional tensor specifying the new shape. 111 When the second input is absent, an extra argument `shape` must be specified. 112 It outputs the reshaped tensor as well as the original shape. 114 At most one dimension of the new shape can be -1. In this case, the value is 115 inferred from the size of the tensor and the remaining dimensions. A dimension 116 could also be 0, in which case the actual dimension value is going to be copied 117 from the input tensor. 119 .Arg("shape",
"New shape")
120 .Input(0,
"data",
"An input tensor.")
121 .Input(1,
"new_shape",
"New shape.")
122 .Output(0,
"reshaped",
"Reshaped data.")
123 .Output(1,
"old_shape",
"Original shape.")
124 .InheritOnnxSchema(
"Reshape");
127 using GradientMakerBase::GradientMakerBase;
128 vector<OperatorDef> GetGradientDefs()
override {
132 vector<string>{GO(0), O(1)},
133 vector<string>{GI(0),
"_" + GI(0) +
"_dims"});
137 bool CopyArguments()
const override {
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
static vector< OperatorDef > SingleGradientDef(const Args &...args)
a helper function to allow one to create one single operator def, which is usually the case for many ...