reshape_op.cc 25.9 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Y
Yibing Liu 已提交
2

L
Luo Tao 已提交
3 4 5
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
Y
Yibing Liu 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Y
Yibing Liu 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
Y
Yibing Liu 已提交
14

Y
Yi Wang 已提交
15 16
#include <string>
#include <vector>
Y
yuyang18 已提交
17
#include "paddle/fluid/framework/op_registry.h"
Y
Yi Wang 已提交
18

Y
Yibing Liu 已提交
19 20 21
namespace paddle {
namespace operators {

22 23 24 25 26 27 28 29
using Tensor = framework::Tensor;

inline std::vector<int> get_new_shape(
    const std::vector<const Tensor *> &list_new_shape_tensor) {
  // get tensor from
  std::vector<int> vec_new_shape;
  for (size_t i = 0; i < list_new_shape_tensor.size(); ++i) {
    auto tensor = list_new_shape_tensor[i];
30 31
    PADDLE_ENFORCE_EQ(
        tensor->dims(), framework::make_ddim({1}),
32 33 34 35 36
        platform::errors::InvalidArgument(
            "If the element type of 'shape' in ReshapeOp is Tensor, "
            "the element's shape must be [1]. But received the element's shape "
            "is [%s]",
            tensor->dims()));
37 38 39 40 41 42 43 44 45 46 47 48 49
    if (platform::is_gpu_place(tensor->place())) {
      framework::Tensor temp;
      TensorCopySync(*tensor, platform::CPUPlace(), &temp);

      vec_new_shape.push_back(static_cast<int32_t>(*temp.data<int32_t>()));
    } else {
      vec_new_shape.push_back(static_cast<int32_t>(*tensor->data<int32_t>()));
    }
  }

  return vec_new_shape;
}

Y
yuyang18 已提交
50 51 52 53 54 55 56 57
class ReshapeOp : public framework::OperatorWithKernel {
 public:
  ReshapeOp(const std::string &type, const framework::VariableNameMap &inputs,
            const framework::VariableNameMap &outputs,
            const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
58
    PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
59 60
                      platform::errors::InvalidArgument(
                          "Input(X) of ReshapeOp should not be null."));
61
    PADDLE_ENFORCE_EQ(ctx->HasOutput("Out"), true,
62 63
                      platform::errors::InvalidArgument(
                          "Output(Out) of ReshapeOp should not be null."));
Y
yuyang18 已提交
64

65 66
    if (ctx->HasInputs("ShapeTensor")) {
      // top prority shape
67
      auto ShapeTensor = ctx->Inputs("ShapeTensor");
68 69
      PADDLE_ENFORCE_GT(
          ShapeTensor.size(), 0,
70 71 72 73 74
          platform::errors::InvalidArgument(
              "When `shape` in ReshapeOp is a list or tuple "
              "which contains Tensor, the shape's size can't be zero. "
              "But received shape's size is %d.",
              ShapeTensor.size()));
75 76 77 78 79 80 81
      auto infer_shape = ctx->Attrs().Get<std::vector<int>>("shape");
      const int64_t copy_dim_val = 0;
      auto in_dims = ctx->GetInputDim("X");
      for (size_t i = 0; i < infer_shape.size(); ++i) {
        if (infer_shape[i] == copy_dim_val) {
          PADDLE_ENFORCE_LT(
              static_cast<int>(i), in_dims.size(),
82 83 84 85 86
              platform::errors::InvalidArgument(
                  "The index of 0 in `shape` must be less than "
                  "the input tensor X's dimensions. But received shape[%d] "
                  "= 0, X's dimensions = %d, X's shape = [%s].",
                  i, in_dims.size(), in_dims));
87 88 89 90 91 92 93
          infer_shape[i] = in_dims[i];
        }
      }
      auto infer_out_dims = framework::make_ddim(infer_shape);
      ctx->SetOutputDim("Out", infer_out_dims);
      return;
    }
Y
yuyang18 已提交
94

95 96 97 98 99 100 101 102 103 104 105
    const std::vector<int> &shape = ctx->Attrs().Get<std::vector<int>>("shape");
    if (ctx->HasInput("Shape") && shape.empty()) {
      auto shape_dims = ctx->GetInputDim("Shape");
      int num_ele = 1;
      for (int i = 0; i < shape_dims.size(); ++i) {
        num_ele *= shape_dims[i];
      }
      auto vec_dims = std::vector<int>(num_ele, -1);
      auto out_dims = framework::make_ddim(vec_dims);
      ctx->SetOutputDim("Out", out_dims);
      ctx->ShareLoD("X", /*->*/ "Out");
106 107
      return;
    }
108 109

    if (ctx->HasInput("Shape") && !shape.empty() && ctx->IsRuntime()) {
Y
yuyang18 已提交
110 111 112 113 114
      // If true, set the shape of Output(Out) according to Input(Shape) in
      // ReshapeKernel with ExecutionContext. Also check LoD in ReshapeKernel.
      ctx->ShareLoD("X", /*->*/ "Out");
      return;
    }
115

116 117 118 119
    PADDLE_ENFORCE_EQ(!shape.empty(), true,
                      platform::errors::InvalidArgument(
                          "The parameter 'shape' in ReshapeOp must be set. "
                          "But received 'shape' is empty."));
Y
yuyang18 已提交
120 121 122 123 124 125 126 127 128 129 130 131 132
    auto x_dims = ctx->GetInputDim("X");
    auto out_dims = ValidateShape(shape, x_dims);
    ctx->SetOutputDim("Out", out_dims);
    if (x_dims[0] == out_dims[0]) {
      // Only pass LoD when the first dimension of output and Input(X)
      // are the same.
      ctx->ShareLoD("X", /*->*/ "Out");
    }
  }

  static framework::DDim ValidateShape(const std::vector<int> shape,
                                       const framework::DDim &in_dims) {
    const int64_t in_size = framework::product(in_dims);
C
chengduo 已提交
133 134 135
    auto in_dims_vec = framework::vectorize(in_dims);
    bool all_positive = std::all_of(in_dims_vec.cbegin(), in_dims_vec.cend(),
                                    [](int64_t i) { return i > 0; });
Y
yuyang18 已提交
136 137 138 139 140 141 142 143 144 145
    // only one dimension can be set to -1, whose size will be automatically
    // infered.
    const int64_t unk_dim_val = -1;
    const int64_t copy_dim_val = 0;

    std::vector<int64_t> output_shape(shape.size(), 0);
    int64_t capacity = 1;
    int unk_dim_idx = -1;
    for (size_t i = 0; i < shape.size(); ++i) {
      if (shape[i] == unk_dim_val) {
146 147
        PADDLE_ENFORCE_EQ(
            unk_dim_idx, -1,
148 149 150 151
            platform::errors::InvalidArgument(
                "Only one dimension value of 'shape' in ReshapeOp can "
                "be -1. But received shape = [%s], shape[%d] is also -1.",
                framework::make_ddim(shape), i));
Y
yuyang18 已提交
152 153
        unk_dim_idx = i;
      } else if (shape[i] == copy_dim_val) {
154 155
        PADDLE_ENFORCE_LT(
            static_cast<int>(i), in_dims.size(),
156 157 158 159 160 161
            platform::errors::InvalidArgument(
                "The index of 0 in `shape` must be less than "
                "the input tensor X's dimensions. "
                "But received shape = [%s], shape[%d] = 0, X's shape = [%s], "
                "X's dimensions = %d.",
                framework::make_ddim(shape), i, in_dims, in_dims.size()));
Y
yuyang18 已提交
162
      } else {
163 164
        PADDLE_ENFORCE_GT(
            shape[i], 0,
165 166
            platform::errors::InvalidArgument(
                "Each dimension value of 'shape' in ReshapeOp must not "
T
tianshuo78520a 已提交
167
                "be negative except one unknown dimension. "
168 169
                "But received  shape = [%s], shape[%d] = %d.",
                framework::make_ddim(shape), i, shape[i]));
Y
yuyang18 已提交
170 171 172 173 174 175 176 177
      }

      capacity *= (shape[i] ? shape[i] : in_dims[i]);
      output_shape[i] =
          (shape[i] ? static_cast<int64_t>(shape[i]) : in_dims[i]);
    }

    if (unk_dim_idx != -1) {
C
chengduo 已提交
178
      if (all_positive) {
Y
yuyang18 已提交
179 180 181 182 183
        // in_size < 0 and is un-determinate in compile time, skip the check,
        // for example, in_dims = [-1, 8, 1, 1], shape = [-1, 3, 8],
        // capacity = -24, in_size = -8, output_shape[0] = 0
        // the following check will fail.
        output_shape[unk_dim_idx] = -in_size / capacity;
184 185 186 187 188 189 190
        PADDLE_ENFORCE_EQ(
            output_shape[unk_dim_idx] * capacity, -in_size,
            platform::errors::InvalidArgument(
                "The 'shape' attribute in ReshapeOp is invalid. "
                "The input tensor X'size must be divisible by known "
                "capacity of 'shape'. "
                "But received X's shape = [%s], X's size = %d, "
191
                "'shape' is [%s], known capacity of 'shape' is %d.",
192
                in_dims, in_size, framework::make_ddim(shape), capacity));
Y
yuyang18 已提交
193 194 195 196
      } else {
        output_shape[unk_dim_idx] = -1;
      }
    } else {
Y
Yamei-Lee 已提交
197 198 199
      if (all_positive) {
        PADDLE_ENFORCE_EQ(
            capacity, in_size,
200 201 202 203 204 205 206
            platform::errors::InvalidArgument(
                "The 'shape' in ReshapeOp is invalid. "
                "The input tensor X'size must be equal to the capacity of "
                "'shape'. "
                "But received X's shape = [%s], X's size = %d, 'shape' is "
                "[%s], the capacity of 'shape' is %d.",
                in_dims, in_size, framework::make_ddim(shape), capacity));
Y
Yamei-Lee 已提交
207
      }
Y
yuyang18 已提交
208 209 210 211 212 213 214
    }
    return framework::make_ddim(output_shape);
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
215 216 217
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "X"),
        ctx.device_context());
Y
yuyang18 已提交
218
  }
219 220 221 222 223 224 225 226 227 228

  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    if (var_name == "ShapeTensor") {
      return expected_kernel_type;
    }
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   tensor.place(), tensor.layout());
  }
Y
yuyang18 已提交
229 230
};

Y
Yibing Liu 已提交
231 232
class ReshapeOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
233
  void Make() override {
234 235
    AddInput("X", "(Tensor). The input tensor of reshape operator.");
    AddInput("Shape",
236 237 238
             "(Tensor<int32>, optional). Target shape of reshape operator. "
             "It has a higher priority than Attr(shape) but a lower priority "
             "than Input(ShapeTensor). The Attr(shape) still should be "
T
tianshuo78520a 已提交
239
             "set correctly to guarantee shape inference in compile time.")
240
        .AsDispensable();
241 242
    AddInput(
        "ShapeTensor",
243 244 245 246
        "(vector<Tensor<int32>>, optional). Target shape of reshape operator. "
        "It has the highest priority compare with Input(Shape) and "
        "Attr(shape)."
        "The shape of the element in vector must be [1].")
247 248
        .AsDuplicable()
        .AsDispensable();
249
    AddOutput("Out", "(Tensor). The output tensor of reshape operator.");
C
caoying03 已提交
250
    AddAttr<std::vector<int>>(
251 252 253 254
        "shape",
        "(std::vector<int>) Target shape of reshape operator."
        "It has the lowest priority compare with Input(Shape) and "
        " Input(ShapeTensor).")
255
        .SetDefault({});
K
kexinzhao 已提交
256 257
    AddComment(R"DOC(
Reshape Operator.
Y
Yibing Liu 已提交
258

259 260
Reshape Input(X) into the shape specified by Attr(shape) or Input(Shape). The
data in Input(X) are unchanged.
Y
Yibing Liu 已提交
261

C
caoying03 已提交
262
Examples:
Y
Yibing Liu 已提交
263

C
caoying03 已提交
264 265 266 267
1. Given a 3-D tensor Input(X) with a shape [2, 4, 6], and the target shape
specified by Attr(shape) is [6, 8], the reshape operator will transform Input(X)
into a 2-D tensor with shape [6, 8] and leaving Input(X)'s data unchanged.

268
2. Given a 3-D tensor Input(X) with a shape [2, 4, 6], and the target shape
C
caoying03 已提交
269 270 271 272 273 274
specified by Attr(shape) is [2, 3, -1, 2], the reshape operator will transform
Input(X) into a 4-D tensor with shape [2, 3, 4, 2] and leaving Input(X)'s data
unchanged. In this case, one and only dimension of Attr(shape) can be set to -1,
the value of this dimension is inferred from the total element number of
Input(X) and remaining dimensions.

275
3. Given a 3-D tensor Input(X) with a shape [2, 4, 6], and the target shape
C
caoying03 已提交
276 277 278 279
specified by Attr(shape) is [-1, 0, 3, 2], the reshape operator will transform
Input(X) into a 4-D tensor with shape [2, 4, 3, 2] and leaving Input(X)'s data
unchanged. In this case, besides -1, 0 means the actual dimension value is going
to be copied from the corresponding dimension of Input(X).
Y
Yibing Liu 已提交
280

C
caoying03 已提交
281
Note:
Y
Yibing Liu 已提交
282

C
caoying03 已提交
283 284 285
1. One and only one dimension in Attr(shape) can be set -1. In this case,
the actual dimension value will be infered from the total element number of
Input(X) and remaining dimensions.
286 287

2. More than one dimensions in Attr(shape) can be set to 0, which means the real
C
caoying03 已提交
288
dimension value will be copied from Input(X) at runtime. Note that the index of
G
guosheng 已提交
289
0 can not exceed Rank(X). For example, Input(X) is a 3-D tensor with shape
C
caoying03 已提交
290
[2, 3, 4], Attr(shape) = [2, 3, 2, 0] is an invalid input.
291 292

3. Input(Shape) has a higher priority than Attr(shape) if it is provided, while
T
tianshuo78520a 已提交
293
Attr(shape) still should be set correctly to guarantee shape inference in
294
compile-time.
Y
Yibing Liu 已提交
295

Y
Yibing Liu 已提交
296 297 298 299 300 301 302 303 304 305 306 307
)DOC");
  }
};

class ReshapeGradOp : public framework::OperatorWithKernel {
 public:
  ReshapeGradOp(const std::string &type,
                const framework::VariableNameMap &inputs,
                const framework::VariableNameMap &outputs,
                const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

308
  void InferShape(framework::InferShapeContext *ctx) const override {
309 310 311
    PADDLE_ENFORCE_EQ(
        ctx->HasInput("X"), true,
        platform::errors::InvalidArgument("Input(X) shouldn't be null."));
312
    PADDLE_ENFORCE_EQ(ctx->HasInput(framework::GradVarName("Out")), true,
313 314
                      platform::errors::InvalidArgument(
                          "Input(Out@GRAD) shouldn't be null."));
Q
Qiao Longfei 已提交
315
    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
Y
Yibing Liu 已提交
316
  }
317 318 319 320

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
321 322 323
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "X"),
        ctx.device_context());
324
  }
Y
Yibing Liu 已提交
325 326
};

Y
yuyang18 已提交
327 328 329 330 331
class ReshapeKernel {
 public:
  void operator()(const framework::ExecutionContext &ctx) const {
    auto *out = ctx.Output<framework::LoDTensor>("Out");
    auto *in = ctx.Input<framework::LoDTensor>("X");
Y
yuyang18 已提交
332

Y
yuyang18 已提交
333
    framework::DDim out_dims = out->dims();
Y
yuyang18 已提交
334

335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357
    auto list_new_shape_tensor =
        ctx.MultiInput<framework::Tensor>("ShapeTensor");
    if (list_new_shape_tensor.size() > 0) {
      // have shape tensor
      auto new_shape = get_new_shape(list_new_shape_tensor);
      out_dims = ReshapeOp::ValidateShape(new_shape, in->dims());

    } else {
      auto *shape_tensor = ctx.HasInput("Shape")
                               ? ctx.Input<framework::LoDTensor>("Shape")
                               : nullptr;

      if (shape_tensor) {
        auto *shape_data = shape_tensor->data<int>();
        framework::Tensor cpu_shape_tensor;
        if (platform::is_gpu_place(shape_tensor->place())) {
          TensorCopySync(*shape_tensor, platform::CPUPlace(),
                         &cpu_shape_tensor);
          shape_data = cpu_shape_tensor.data<int>();
        }
        auto shape =
            std::vector<int>(shape_data, shape_data + shape_tensor->numel());
        out_dims = ReshapeOp::ValidateShape(shape, in->dims());
Y
yuyang18 已提交
358 359
      }
    }
Y
yuyang18 已提交
360

361
    out->Resize(out_dims);
362
    out->mutable_data(ctx.GetPlace(), in->type());
Y
Yiqun Liu 已提交
363 364 365
    framework::TensorCopy(
        *in, ctx.GetPlace(),
        ctx.template device_context<platform::DeviceContext>(), out);
Y
yuyang18 已提交
366 367
    out->Resize(out_dims);
  }
Y
yuyang18 已提交
368 369 370 371 372 373 374
};

class ReshapeGradKernel {
 public:
  void operator()(const framework::ExecutionContext &ctx) const {
    auto *d_out = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
    auto *d_x = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
D
dzhwinter 已提交
375
    auto in_dims = d_x->dims();
Y
yuyang18 已提交
376

377 378
    d_x->mutable_data(ctx.GetPlace(), d_out->type());
    framework::TensorCopySync(*d_out, ctx.GetPlace(), d_x);
D
dzhwinter 已提交
379
    d_x->Resize(in_dims);
Y
yuyang18 已提交
380
  }
Y
yuyang18 已提交
381 382
};

383 384 385 386 387 388 389 390 391 392 393 394 395 396
class ReshapeDoubleGradKernel {
 public:
  void operator()(const framework::ExecutionContext &ctx) const {
    auto *dd_x = ctx.Input<framework::Tensor>("DDX");
    auto *dd_out = ctx.Output<framework::Tensor>("DDOut");

    auto out_dims = dd_out->dims();

    dd_out->mutable_data(ctx.GetPlace(), dd_x->type());
    framework::TensorCopySync(*dd_x, ctx.GetPlace(), dd_out);
    dd_out->Resize(out_dims);
  }
};

397 398 399 400 401 402 403 404 405 406 407 408 409
// FIXME(zcd): reshape2 adds an intermediate output(XShape) based on reshape,
// the XShape is used to carry the shape and lod of X which will be used in
// reshape_grad, in this way, the framework can reuse the memory of X
// immediately the reshape_op is finished.
// Considering compatibility issues, we could not fix reshape_op
class Reshape2Op : public ReshapeOp {
 public:
  Reshape2Op(const std::string &type, const framework::VariableNameMap &inputs,
             const framework::VariableNameMap &outputs,
             const framework::AttributeMap &attrs)
      : ReshapeOp(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
410
    PADDLE_ENFORCE_EQ(ctx->HasOutput("XShape"), true,
411 412
                      platform::errors::InvalidArgument(
                          "Output(XShape) of ReshapeOp should not be null."));
413 414 415 416 417 418 419 420
    const auto &x_dims = ctx->GetInputDim("X");
    std::vector<int64_t> xshape_dims(x_dims.size() + 1);
    xshape_dims[0] = 0;
    for (int i = 0; i < x_dims.size(); ++i) {
      xshape_dims[i + 1] = x_dims[i];
    }
    ctx->SetOutputDim("XShape", framework::make_ddim(xshape_dims));
    ctx->ShareLoD("X", /*->*/ "XShape");
M
minqiyang 已提交
421 422

    ReshapeOp::InferShape(ctx);
423 424 425 426 427 428 429 430 431 432 433
  }
};

class Reshape2OpMaker : public ReshapeOpMaker {
 public:
  void Make() override {
    ReshapeOpMaker::Make();
    AddOutput("XShape",
              "XShape is just used to store the shape and lod of X, which will "
              "be used in FlattenGradOp.")
        .AsIntermediate();
434 435 436 437
    AddAttr<bool>(
        "use_quantizer",
        "(bool, default false) "
        "This parameter is no longer used. Use 'mkldnn_data_type' instead.")
438
        .SetDefault(false);
439 440 441 442 443
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
        .InEnum({"float32", "int8", "bfloat16"});
444 445 446
  }
};

H
hong 已提交
447 448
template <typename T>
class Reshape2GradMaker : public framework::SingleGradOpMaker<T> {
449
 public:
H
hong 已提交
450
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
451

452
  void Apply(GradOpPtr<T> grad_op) const override {
453
    grad_op->SetType("reshape2_grad");
H
hong 已提交
454 455 456 457
    grad_op->SetInput("XShape", this->Output("XShape"));
    grad_op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
    grad_op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
    grad_op->SetAttrMap(this->Attrs());
458 459 460
  }
};

H
hong 已提交
461 462
template <typename T>
class Reshape2DoubleGradMaker : public framework::SingleGradOpMaker<T> {
463
 public:
H
hong 已提交
464
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
465

466
  void Apply(GradOpPtr<T> grad_op) const override {
467
    grad_op->SetType("reshape2_grad_grad");
H
hong 已提交
468 469 470 471
    grad_op->SetInput("DOut", this->Input(framework::GradVarName("Out")));
    grad_op->SetInput("DDX", this->OutputGrad(framework::GradVarName("X")));
    grad_op->SetOutput("DDOut", this->InputGrad(framework::GradVarName("Out")));
    grad_op->SetAttrMap(this->Attrs());
472 473 474
  }
};

475 476 477 478 479 480 481 482 483
class Reshape2GradOp : public framework::OperatorWithKernel {
 public:
  Reshape2GradOp(const std::string &type,
                 const framework::VariableNameMap &inputs,
                 const framework::VariableNameMap &outputs,
                 const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
484 485 486
    PADDLE_ENFORCE_EQ(
        ctx->HasInput("XShape"), true,
        platform::errors::InvalidArgument("Input(XShape) shouldn't be null."));
487
    PADDLE_ENFORCE_EQ(ctx->HasInput(framework::GradVarName("Out")), true,
488 489
                      platform::errors::InvalidArgument(
                          "Input(Out@GRAD) shouldn't be null."));
490 491 492 493 494 495 496 497 498
    auto xshape_dims = ctx->GetInputDim("XShape");
    auto x_dims = framework::slice_ddim(xshape_dims, 1, xshape_dims.size());
    ctx->SetOutputDim(framework::GradVarName("X"), x_dims);
    ctx->ShareLoD("XShape", framework::GradVarName("X"));
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
499 500 501
    return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(
                                       ctx, framework::GradVarName("Out")),
                                   ctx.device_context());
502
  }
503 504 505 506 507 508 509 510 511 512

  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    if (var_name == "ShapeTensor") {
      return expected_kernel_type;
    }
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   tensor.place(), tensor.layout());
  }
513 514
};

515 516 517 518 519 520 521 522 523 524
class Reshape2DoubleGradOp : public framework::OperatorWithKernel {
 public:
  Reshape2DoubleGradOp(const std::string &type,
                       const framework::VariableNameMap &inputs,
                       const framework::VariableNameMap &outputs,
                       const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
    PADDLE_ENFORCE_EQ(ctx->HasInput("DDX"), true,
525 526
                      platform::errors::InvalidArgument(
                          "Input(X@GRAD_GRAD) shouldn't be null."));
527 528 529 530 531 532 533 534
    if (ctx->HasOutput("DDOut") && ctx->HasInput("DDX")) {
      ctx->ShareDim("DOut", "DDOut");
    }
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
535 536 537
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "DDX"),
        ctx.device_context());
538 539 540 541 542 543 544 545 546 547 548 549 550
  }

  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    if (var_name == "ShapeTensor") {
      return expected_kernel_type;
    }
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   tensor.place(), tensor.layout());
  }
};

551 552
DECLARE_INPLACE_OP_INFERER(ReshapeOpInplaceInferer, {"X", "Out"});
DECLARE_INPLACE_OP_INFERER(ReshapeGradInplaceInferer,
553 554
                           {framework::GradVarName("Out"),
                            framework::GradVarName("X")});
555 556
DECLARE_INPLACE_OP_INFERER(ReshapeDoubleGradInplaceInferer, {"DDX", "DDOut"});
DECLARE_NO_NEED_BUFFER_VARS_INFERER(ReshapeDoubleGradOpNoNeedBufferVarInferer,
Z
Zeng Jinle 已提交
557
                                    "DOut");
D
dzhwinter 已提交
558

Y
Yibing Liu 已提交
559 560 561
}  // namespace operators
}  // namespace paddle
namespace ops = paddle::operators;
562
namespace plat = paddle::platform;
Y
Yibing Liu 已提交
563

H
hong 已提交
564 565 566 567
REGISTER_OPERATOR(
    reshape, ops::ReshapeOp, ops::ReshapeOpMaker,
    paddle::framework::DefaultGradOpMaker<paddle::framework::OpDesc, true>,
    paddle::framework::DefaultGradOpMaker<paddle::imperative::OpBase, true>,
568
    ops::ReshapeOpInplaceInferer);
D
dzhwinter 已提交
569
REGISTER_OPERATOR(reshape_grad, ops::ReshapeGradOp,
570
                  ops::ReshapeGradInplaceInferer);
571

572 573 574 575 576 577 578
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape, float, ops::ReshapeKernel, double,
                               ops::ReshapeKernel, int, ops::ReshapeKernel,
                               int64_t, ops::ReshapeKernel);
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape_grad, float, ops::ReshapeGradKernel,
                               double, ops::ReshapeGradKernel, int,
                               ops::ReshapeGradKernel, int64_t,
                               ops::ReshapeGradKernel);
579
REGISTER_OPERATOR(reshape2, ops::Reshape2Op, ops::Reshape2OpMaker,
H
hong 已提交
580 581
                  ops::Reshape2GradMaker<paddle::framework::OpDesc>,
                  ops::Reshape2GradMaker<paddle::imperative::OpBase>,
582
                  ops::ReshapeOpInplaceInferer);
D
dzhwinter 已提交
583
REGISTER_OPERATOR(reshape2_grad, ops::Reshape2GradOp,
H
hong 已提交
584 585
                  ops::Reshape2DoubleGradMaker<paddle::framework::OpDesc>,
                  ops::Reshape2DoubleGradMaker<paddle::imperative::OpBase>,
586
                  ops::ReshapeGradInplaceInferer);
587
REGISTER_OPERATOR(reshape2_grad_grad, ops::Reshape2DoubleGradOp,
588 589
                  ops::ReshapeDoubleGradInplaceInferer,
                  ops::ReshapeDoubleGradOpNoNeedBufferVarInferer);
590

591
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2, float, ops::ReshapeKernel, double,
592 593 594
                               ops::ReshapeKernel, int8_t, ops::ReshapeKernel,
                               uint8_t, ops::ReshapeKernel, int,
                               ops::ReshapeKernel, int64_t, ops::ReshapeKernel);
595 596 597 598
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2_grad, float, ops::ReshapeGradKernel,
                               double, ops::ReshapeGradKernel, int,
                               ops::ReshapeGradKernel, int64_t,
                               ops::ReshapeGradKernel);
599 600 601 602 603
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2_grad_grad, float,
                               ops::ReshapeDoubleGradKernel, double,
                               ops::ReshapeDoubleGradKernel, int,
                               ops::ReshapeDoubleGradKernel, int64_t,
                               ops::ReshapeDoubleGradKernel);
604

Y
yuyang18 已提交
605
#ifdef PADDLE_WITH_CUDA
606 607
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape, float, ops::ReshapeKernel, double,
                                ops::ReshapeKernel, int, ops::ReshapeKernel,
608 609
                                int64_t, ops::ReshapeKernel, plat::float16,
                                ops::ReshapeKernel);
610 611 612
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape_grad, float, ops::ReshapeGradKernel,
                                double, ops::ReshapeGradKernel, int,
                                ops::ReshapeGradKernel, int64_t,
613
                                ops::ReshapeGradKernel, plat::float16,
614 615 616
                                ops::ReshapeGradKernel);
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape2, float, ops::ReshapeKernel, double,
                                ops::ReshapeKernel, int, ops::ReshapeKernel,
617 618
                                int64_t, ops::ReshapeKernel, plat::float16,
                                ops::ReshapeKernel);
619 620 621
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape2_grad, float, ops::ReshapeGradKernel,
                                double, ops::ReshapeGradKernel, int,
                                ops::ReshapeGradKernel, int64_t,
622
                                ops::ReshapeGradKernel, plat::float16,
623
                                ops::ReshapeGradKernel);
624 625 626 627 628 629 630

REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape2_grad_grad, float,
                                ops::ReshapeDoubleGradKernel, double,
                                ops::ReshapeDoubleGradKernel, int,
                                ops::ReshapeDoubleGradKernel, int64_t,
                                ops::ReshapeDoubleGradKernel, plat::float16,
                                ops::ReshapeDoubleGradKernel);
Y
yuyang18 已提交
631
#endif