reshape_op.cc 24.7 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 32 33 34 35
    PADDLE_ENFORCE_EQ(
        tensor->dims(), framework::make_ddim({1}),
        "ShapeError: 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());
36 37 38 39 40 41 42 43 44 45 46 47 48
    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 已提交
49 50 51 52 53 54 55 56
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 {
57 58 59 60
    PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
                      "Input(X) of ReshapeOp should not be null.");
    PADDLE_ENFORCE_EQ(ctx->HasOutput("Out"), true,
                      "Output(Out) of ReshapeOp should not be null.");
Y
yuyang18 已提交
61

62 63
    if (ctx->HasInputs("ShapeTensor")) {
      // top prority shape
64
      auto ShapeTensor = ctx->Inputs("ShapeTensor");
65 66 67 68 69 70
      PADDLE_ENFORCE_GT(
          ShapeTensor.size(), 0,
          "ShapeError: 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());
71 72 73 74 75 76 77
      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(),
78 79 80 81
              "ShapeError: 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);
82 83 84 85 86 87 88
          infer_shape[i] = in_dims[i];
        }
      }
      auto infer_out_dims = framework::make_ddim(infer_shape);
      ctx->SetOutputDim("Out", infer_out_dims);
      return;
    }
Y
yuyang18 已提交
89

90 91 92 93 94 95 96 97 98 99 100
    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");
101 102
      return;
    }
103 104

    if (ctx->HasInput("Shape") && !shape.empty() && ctx->IsRuntime()) {
Y
yuyang18 已提交
105 106 107 108 109
      // 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;
    }
110

111 112 113 114
    PADDLE_ENFORCE_EQ(
        !shape.empty(), true,
        "ShapeError: The parameter 'shape' in ReshapeOp must be set. "
        "But received 'shape' is empty.");
Y
yuyang18 已提交
115 116 117 118 119 120 121 122 123 124 125 126 127
    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 已提交
128 129 130
    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 已提交
131 132 133 134 135 136 137 138 139 140
    // 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) {
141 142
        PADDLE_ENFORCE_EQ(
            unk_dim_idx, -1,
143 144 145
            "ShapeError: 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 已提交
146 147
        unk_dim_idx = i;
      } else if (shape[i] == copy_dim_val) {
148 149
        PADDLE_ENFORCE_LT(
            static_cast<int>(i), in_dims.size(),
150 151 152 153 154
            "ShapeError: 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 已提交
155
      } else {
156 157
        PADDLE_ENFORCE_GT(
            shape[i], 0,
158 159 160 161
            "ShapeError: Each dimension value of 'shape' in ReshapeOp must not "
            "be negtive except one unknown dimension. "
            "But received  shape = [%s], shape[%d] = %d.",
            framework::make_ddim(shape), i, shape[i]);
Y
yuyang18 已提交
162 163 164 165 166 167 168 169
      }

      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 已提交
170
      if (all_positive) {
Y
yuyang18 已提交
171 172 173 174 175 176
        // 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;
        PADDLE_ENFORCE_EQ(output_shape[unk_dim_idx] * capacity, -in_size,
177 178 179 180 181 182 183 184
                          "ShapeError: The 'shape' 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, "
                          "'shape' is [%s], known "
                          "capacity of 'shape' is %d.",
                          in_dims, in_size, framework::make_ddim(shape),
                          capacity);
Y
yuyang18 已提交
185 186 187 188
      } else {
        output_shape[unk_dim_idx] = -1;
      }
    } else {
Y
Yamei-Lee 已提交
189 190 191 192 193 194 195 196 197 198
      if (all_positive) {
        PADDLE_ENFORCE_EQ(
            capacity, in_size,
            "ShapeError: 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
yuyang18 已提交
199 200 201 202 203 204 205
    }
    return framework::make_ddim(output_shape);
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
206 207 208
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "X"),
        ctx.device_context());
Y
yuyang18 已提交
209
  }
210 211 212 213 214 215 216 217 218 219

  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 已提交
220 221
};

Y
Yibing Liu 已提交
222 223
class ReshapeOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
224
  void Make() override {
225 226
    AddInput("X", "(Tensor). The input tensor of reshape operator.");
    AddInput("Shape",
227 228 229
             "(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 "
230 231
             "set correctly to gurantee shape inference in compile time.")
        .AsDispensable();
232 233
    AddInput(
        "ShapeTensor",
234 235 236 237
        "(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].")
238 239
        .AsDuplicable()
        .AsDispensable();
240
    AddOutput("Out", "(Tensor). The output tensor of reshape operator.");
C
caoying03 已提交
241
    AddAttr<std::vector<int>>(
242 243 244 245
        "shape",
        "(std::vector<int>) Target shape of reshape operator."
        "It has the lowest priority compare with Input(Shape) and "
        " Input(ShapeTensor).")
246
        .SetDefault({});
K
kexinzhao 已提交
247 248
    AddComment(R"DOC(
Reshape Operator.
Y
Yibing Liu 已提交
249

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

C
caoying03 已提交
253
Examples:
Y
Yibing Liu 已提交
254

C
caoying03 已提交
255 256 257 258
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.

259
2. Given a 3-D tensor Input(X) with a shape [2, 4, 6], and the target shape
C
caoying03 已提交
260 261 262 263 264 265
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.

266
3. Given a 3-D tensor Input(X) with a shape [2, 4, 6], and the target shape
C
caoying03 已提交
267 268 269 270
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 已提交
271

C
caoying03 已提交
272
Note:
Y
Yibing Liu 已提交
273

C
caoying03 已提交
274 275 276
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.
277 278

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

3. Input(Shape) has a higher priority than Attr(shape) if it is provided, while
M
minqiyang 已提交
284
Attr(shape) still should be set correctly to gurantee shape inference in
285
compile-time.
Y
Yibing Liu 已提交
286

Y
Yibing Liu 已提交
287 288 289 290 291 292 293 294 295 296 297 298
)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) {}

299
  void InferShape(framework::InferShapeContext *ctx) const override {
300 301 302
    PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true, "Input(X) shouldn't be null.");
    PADDLE_ENFORCE_EQ(ctx->HasInput(framework::GradVarName("Out")), true,
                      "Input(Out@GRAD) shouldn't be null.");
Q
Qiao Longfei 已提交
303
    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
Y
Yibing Liu 已提交
304
  }
305 306 307 308

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
309 310 311
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "X"),
        ctx.device_context());
312
  }
Y
Yibing Liu 已提交
313 314
};

Y
yuyang18 已提交
315 316 317 318 319
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 已提交
320

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

323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345
    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 已提交
346 347
      }
    }
Y
yuyang18 已提交
348

349
    out->Resize(out_dims);
350
    out->mutable_data(ctx.GetPlace(), in->type());
Y
Yiqun Liu 已提交
351 352 353
    framework::TensorCopy(
        *in, ctx.GetPlace(),
        ctx.template device_context<platform::DeviceContext>(), out);
Y
yuyang18 已提交
354 355
    out->Resize(out_dims);
  }
Y
yuyang18 已提交
356 357 358 359 360 361 362
};

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 已提交
363
    auto in_dims = d_x->dims();
Y
yuyang18 已提交
364

365 366
    d_x->mutable_data(ctx.GetPlace(), d_out->type());
    framework::TensorCopySync(*d_out, ctx.GetPlace(), d_x);
D
dzhwinter 已提交
367
    d_x->Resize(in_dims);
Y
yuyang18 已提交
368
  }
Y
yuyang18 已提交
369 370
};

371 372 373 374 375 376 377 378 379 380 381 382 383 384
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);
  }
};

385 386 387 388 389 390 391 392 393 394 395 396 397
// 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 {
398 399
    PADDLE_ENFORCE_EQ(ctx->HasOutput("XShape"), true,
                      "Output(XShape) of ReshapeOp should not be null.");
400 401 402 403 404 405 406 407
    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 已提交
408 409

    ReshapeOp::InferShape(ctx);
410 411 412 413 414 415 416 417 418 419 420 421 422 423
  }
};

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();
  }
};

H
hong 已提交
424 425
template <typename T>
class Reshape2GradMaker : public framework::SingleGradOpMaker<T> {
426
 public:
H
hong 已提交
427
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
428

H
hong 已提交
429 430
  std::unique_ptr<T> Apply() const override {
    auto *grad_op = new T();
431
    grad_op->SetType("reshape2_grad");
H
hong 已提交
432 433 434 435 436 437
    grad_op->SetInput("XShape", this->Output("XShape"));
    grad_op->SetInput("ShapeTensor", this->Input("ShapeTensor"));
    grad_op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
    grad_op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
    grad_op->SetAttrMap(this->Attrs());
    return std::unique_ptr<T>(grad_op);
438 439 440
  }
};

H
hong 已提交
441 442
template <typename T>
class Reshape2DoubleGradMaker : public framework::SingleGradOpMaker<T> {
443
 public:
H
hong 已提交
444
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
445

H
hong 已提交
446 447
  std::unique_ptr<T> Apply() const override {
    auto *grad_op = new T();
448 449
    grad_op->SetType("reshape2_grad_grad");

H
hong 已提交
450 451 452
    grad_op->SetInput("ShapeTensor", this->Input("ShapeTensor"));
    grad_op->SetInput("DOut", this->Input(framework::GradVarName("Out")));
    grad_op->SetInput("DDX", this->OutputGrad(framework::GradVarName("X")));
453

H
hong 已提交
454 455 456
    grad_op->SetOutput("DDOut", this->InputGrad(framework::GradVarName("Out")));
    grad_op->SetAttrMap(this->Attrs());
    return std::unique_ptr<T>(grad_op);
457 458 459
  }
};

460 461 462 463 464 465 466 467 468
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 {
469 470 471 472
    PADDLE_ENFORCE_EQ(ctx->HasInput("XShape"), true,
                      "Input(XShape) shouldn't be null.");
    PADDLE_ENFORCE_EQ(ctx->HasInput(framework::GradVarName("Out")), true,
                      "Input(Out@GRAD) shouldn't be null.");
473 474 475 476 477 478 479 480 481
    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 {
482 483 484
    return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(
                                       ctx, framework::GradVarName("Out")),
                                   ctx.device_context());
485
  }
486 487 488 489 490 491 492 493 494 495

  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());
  }
496 497
};

498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517
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,
                      "Input(X@GRAD_GRAD) shouldn't be null.");

    if (ctx->HasOutput("DDOut") && ctx->HasInput("DDX")) {
      ctx->ShareDim("DOut", "DDOut");
    }
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
518 519 520
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "DDX"),
        ctx.device_context());
521 522 523 524 525 526 527 528 529 530 531 532 533
  }

  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());
  }
};

534 535 536 537
DECLARE_INPLACE_OP_INFERER(ReshapeOpInplaceInToOut, {"X", "Out"});
DECLARE_INPLACE_OP_INFERER(ReshapeGradInplaceInToOut,
                           {framework::GradVarName("Out"),
                            framework::GradVarName("X")});
538
DECLARE_INPLACE_OP_INFERER(ReshapeDoubleGradInplaceInToOut, {"DDX", "DDOut"});
D
dzhwinter 已提交
539

Y
Yibing Liu 已提交
540 541 542
}  // namespace operators
}  // namespace paddle
namespace ops = paddle::operators;
543
namespace plat = paddle::platform;
Y
Yibing Liu 已提交
544

H
hong 已提交
545 546 547 548 549
REGISTER_OPERATOR(
    reshape, ops::ReshapeOp, ops::ReshapeOpMaker,
    paddle::framework::DefaultGradOpMaker<paddle::framework::OpDesc, true>,
    paddle::framework::DefaultGradOpMaker<paddle::imperative::OpBase, true>,
    ops::ReshapeOpInplaceInToOut);
D
dzhwinter 已提交
550 551
REGISTER_OPERATOR(reshape_grad, ops::ReshapeGradOp,
                  ops::ReshapeGradInplaceInToOut);
552

553 554 555 556 557 558 559
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);
560
REGISTER_OPERATOR(reshape2, ops::Reshape2Op, ops::Reshape2OpMaker,
H
hong 已提交
561 562 563
                  ops::Reshape2GradMaker<paddle::framework::OpDesc>,
                  ops::Reshape2GradMaker<paddle::imperative::OpBase>,
                  ops::ReshapeOpInplaceInToOut);
D
dzhwinter 已提交
564
REGISTER_OPERATOR(reshape2_grad, ops::Reshape2GradOp,
H
hong 已提交
565 566 567
                  ops::Reshape2DoubleGradMaker<paddle::framework::OpDesc>,
                  ops::Reshape2DoubleGradMaker<paddle::imperative::OpBase>,
                  ops::ReshapeGradInplaceInToOut);
568 569 570
REGISTER_OPERATOR(reshape2_grad_grad, ops::Reshape2DoubleGradOp,
                  ops::ReshapeDoubleGradInplaceInToOut);

571 572 573 574 575 576 577
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2, float, ops::ReshapeKernel, double,
                               ops::ReshapeKernel, int, ops::ReshapeKernel,
                               int64_t, ops::ReshapeKernel);
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2_grad, float, ops::ReshapeGradKernel,
                               double, ops::ReshapeGradKernel, int,
                               ops::ReshapeGradKernel, int64_t,
                               ops::ReshapeGradKernel);
578 579 580 581 582
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2_grad_grad, float,
                               ops::ReshapeDoubleGradKernel, double,
                               ops::ReshapeDoubleGradKernel, int,
                               ops::ReshapeDoubleGradKernel, int64_t,
                               ops::ReshapeDoubleGradKernel);
583

Y
yuyang18 已提交
584
#ifdef PADDLE_WITH_CUDA
585 586
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape, float, ops::ReshapeKernel, double,
                                ops::ReshapeKernel, int, ops::ReshapeKernel,
587 588
                                int64_t, ops::ReshapeKernel, plat::float16,
                                ops::ReshapeKernel);
589 590 591
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape_grad, float, ops::ReshapeGradKernel,
                                double, ops::ReshapeGradKernel, int,
                                ops::ReshapeGradKernel, int64_t,
592
                                ops::ReshapeGradKernel, plat::float16,
593 594 595
                                ops::ReshapeGradKernel);
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape2, float, ops::ReshapeKernel, double,
                                ops::ReshapeKernel, int, ops::ReshapeKernel,
596 597
                                int64_t, ops::ReshapeKernel, plat::float16,
                                ops::ReshapeKernel);
598 599 600
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape2_grad, float, ops::ReshapeGradKernel,
                                double, ops::ReshapeGradKernel, int,
                                ops::ReshapeGradKernel, int64_t,
601
                                ops::ReshapeGradKernel, plat::float16,
602
                                ops::ReshapeGradKernel);
603 604 605 606 607 608 609

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 已提交
610
#endif