Cherry pick bug fix for Ops: reshape,concat, split and squeeze (#20929)

* [cherry-pick]fix bug in reshape: (#20781)

consider the situation that shape of input can contain more than one -1.

* [cherry-pick]support Tensor for split and concat, support -1 in num_or_sections, add check num_or_sections (#20780)

* improve split and concat op:
1. support Tensor for argument 'dim' in split op.
2. support Tensor for argument 'axis' in concat op.
* redefine function GetDataFromTensor and set unknown output shape to - 1.
* add check: Attr(sections) match Input(X).
* support Tensor for attr(sections) and attr(sections) can contain -1.
* modify error message and fix bug for concat and call Resize only when necessary.
test=release/1.6

* [cherry-pick]improve unsqueeze op to support int, Tensor for argument axes (#20824)

* improve unsqueeze op to support int, Tensor and Tensor list for argument axes.
* call Resize only when necessary. test=release/1.6

* [cherry-pick]Compatible int32 and int64 for attr in concat/split/unsqueeze. test=release/1.6 (#20912)

Author: liym27

paddle/fluid/operators/concat_op.cc

@@ -32,58 +32,36 @@ class ConcatOp : public framework::OperatorWithKernel {
   void InferShape(framework::InferShapeContext *ctx) const override {
     PADDLE_ENFORCE_GE(ctx->Inputs("X").size(), 1UL,
                       "Inputs(X) of ConcatOp should not be empty.");
-    PADDLE_ENFORCE(ctx->HasOutput("Out"),
-                   "Output(Out) of ConcatOp should not be null.");
+    PADDLE_ENFORCE_EQ(ctx->HasOutput("Out"), true,
+                      "Output(Out) of ConcatOp should not be null.");
 
-    auto ins = ctx->GetInputsDim("X");
-    size_t axis =
-        ComputeAxis(static_cast<int64_t>(ctx->Attrs().Get<int>("axis")),
-                    static_cast<int64_t>(ins[0].size()));
+    auto inputs_dims = ctx->GetInputsDim("X");
 
-    const size_t n = ins.size();
-    PADDLE_ENFORCE_GT(n, 0,
+    const size_t inputs_num = inputs_dims.size();
+    PADDLE_ENFORCE_GT(inputs_num, 0,
                       "ShapeError: Input tensors count should > 0. But "
                       "recevied inputs' length is 0.");
-    if (n == 1) {
+    if (inputs_num == 1) {
       VLOG(3) << "Warning: concat op have only one input, may waste memory";
     }
 
-    auto out_dims = ins[0];
-    size_t in_zero_dims_size = out_dims.size();
-    for (size_t i = 1; i < n; i++) {
-      for (size_t j = 0; j < in_zero_dims_size; j++) {
-        if (j == axis) {
-          if (ctx->IsRuntime()) {
-            out_dims[axis] += ins[i][j];
-          } else {
-            if (ins[i][j] == -1) {
-              out_dims[axis] = -1;
-            } else {
-              out_dims[axis] += ins[i][j];
-            }
-          }
-        } else {
-          bool check_shape =
-              ctx->IsRuntime() || (out_dims[j] > 0 && ins[i][j] > 0);
-          if (check_shape) {
-            // check all shape in run time
-            PADDLE_ENFORCE_EQ(
-                out_dims[j], ins[i][j],
-                "ShapeError: Input tensors should have same "
-                "dimensions(or specific dimension = -1) except the axis. "
-                "But recevied axis = %s, input[0]'s shape = "
-                "[%s], input[%s]'s shape = [%s], the \"%s\" "
-                "dimension of input[%s] is unexpected",
-                axis, ins[0], i, ins[j], j, i);
-          }
-        }
+    if (ctx->HasInput("AxisTensor")) {
+      auto out_dims =
+          framework::make_ddim(std::vector<int>(inputs_dims[0].size(), -1));
+      ctx->SetOutputDim("Out", out_dims);
+      ctx->ShareLoD("X", /*->*/ "Out");
+    } else {
+      size_t axis =
+          ComputeAxis(static_cast<int64_t>(ctx->Attrs().Get<int>("axis")),
+                      static_cast<int64_t>(inputs_dims[0].size()));
+      framework::DDim out_dims =
+          ComputeAndCheckShape(ctx->IsRuntime(), inputs_dims, axis);
+      if (out_dims[axis] < 0) {
+        out_dims[axis] = -1;
       }
+      ctx->SetOutputDim("Out", out_dims);
+      ctx->ShareLoD("X", /*->*/ "Out");
     }
-    if (out_dims[axis] < 0) {
-      out_dims[axis] = -1;
-    }
-    ctx->SetOutputDim("Out", out_dims);
-    ctx->ShareLoD("X", /*->*/ "Out");
   }
 
  protected:
@@ -111,6 +89,16 @@ class ConcatOp : public framework::OperatorWithKernel {
 #endif
     return framework::OpKernelType(input_data_type, ctx.GetPlace());
   }
+
+  framework::OpKernelType GetKernelTypeForVar(
+      const std::string &var_name, const Tensor &tensor,
+      const framework::OpKernelType &expected_kernel_type) const override {
+    if (var_name == "AxisTensor") {
+      return expected_kernel_type;
+    }
+    return framework::OpKernelType(expected_kernel_type.data_type_,
+                                   tensor.place(), tensor.layout());
+  }
 };
 
 class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
@@ -128,6 +116,12 @@ class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
                  "interpreted as counting from the end of the rank."
                  "i.e., axis + rank(X) th dimension.")
         .SetDefault(0);
+    AddInput("AxisTensor",
+             "(Tensor) The axis along which the input tensors will be "
+             "concatenated.  "
+             "It has higher priority than Attr(axis). "
+             "The shape of AxisTensor must be [1].")
+        .AsDispensable();
     AddAttr<bool>("use_quantizer",
                   "(bool, default false) "
                   "Set to true for operators that should be quantized and use "
@@ -178,6 +172,16 @@ class ConcatOpGrad : public framework::OperatorWithKernel {
                                        ctx, framework::GradVarName("Out")),
                                    ctx.GetPlace());
   }
+
+  framework::OpKernelType GetKernelTypeForVar(
+      const std::string &var_name, const Tensor &tensor,
+      const framework::OpKernelType &expected_kernel_type) const override {
+    if (var_name == "AxisTensor") {
+      return expected_kernel_type;
+    }
+    return framework::OpKernelType(expected_kernel_type.data_type_,
+                                   tensor.place(), tensor.layout());
+  }
 };
 
 DECLARE_NO_NEED_BUFFER_VARS_INFERENCE(ConcatOpGradNoNeedBufferVarInference,
@@ -192,6 +196,7 @@ class ConcatGradOpDescMaker : public framework::SingleGradOpDescMaker {
     std::unique_ptr<framework::OpDesc> op(new framework::OpDesc());
     op->SetType("concat_grad");
     op->SetInput("X", Input("X"));
+    op->SetInput("AxisTensor", Input("AxisTensor"));
     op->SetInput(framework::GradVarName("Out"), OutputGrad("Out"));
     op->SetOutput(framework::GradVarName("X"), InputGrad("X", false));
     op->SetAttrMap(Attrs());

paddle/fluid/operators/concat_op.h

@@ -14,14 +14,51 @@ limitations under the License. */
 
 #pragma once
 
+#include <string>
 #include <utility>
 #include <vector>
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/operators/math/concat_and_split.h"
 #include "paddle/fluid/operators/strided_memcpy.h"
+#include "paddle/fluid/operators/utils.h"
 
 namespace paddle {
 namespace operators {
+static inline framework::DDim ComputeAndCheckShape(
+    const bool is_runtime, const std::vector<framework::DDim>& inputs_dims,
+    const int axis) {
+  const size_t n = inputs_dims.size();
+  auto out_dims = inputs_dims[0];
+  size_t in_zero_dims_size = out_dims.size();
+  for (size_t i = 1; i < n; i++) {
+    for (size_t j = 0; j < in_zero_dims_size; j++) {
+      if (j == axis) {
+        if (is_runtime) {
+          out_dims[axis] += inputs_dims[i][j];
+        } else {
+          if (inputs_dims[i][j] == -1) {
+            out_dims[axis] = -1;
+          } else {
+            out_dims[axis] += inputs_dims[i][j];
+          }
+        }
+      } else {
+        bool check_shape =
+            is_runtime || (out_dims[j] > 0 && inputs_dims[i][j] > 0);
+        if (check_shape) {
+          // check all shape in run time
+          PADDLE_ENFORCE_EQ(
+              inputs_dims[0][j], inputs_dims[i][j],
+              "ShapeError: Dimension %d in inputs' shapes must be equal. "
+              "But recevied input[0]'s shape = "
+              "[%s], input[%d]'s shape = [%s].",
+              j, inputs_dims[0], i, inputs_dims[i]);
+        }
+      }
+    }
+  }
+  return out_dims;
+}
 
 static inline int64_t ComputeAxis(int64_t axis, int64_t rank) {
   if (axis < 0) {
@@ -36,9 +73,27 @@ class ConcatKernel : public framework::OpKernel<T> {
   void Compute(const framework::ExecutionContext& ctx) const override {
     auto ins = ctx.MultiInput<framework::Tensor>("X");
     framework::Tensor* out = ctx.Output<framework::Tensor>("Out");
-    PADDLE_ENFORCE(ins[0], "The input should not be null.");
-    auto axis = ComputeAxis(static_cast<int64_t>(ctx.Attr<int>("axis")),
-                            static_cast<int64_t>(ins[0]->dims().size()));
+    PADDLE_ENFORCE_EQ(ins[0] != nullptr, true, "The input should not be null.");
+    auto axis = ctx.Attr<int>("axis");
+    bool need_resize_out_dims = false;
+    if (ctx.HasInput("AxisTensor")) {
+      auto* axis_tensor = ctx.Input<framework::Tensor>("AxisTensor");
+      axis = GetDataFromTensor<int>(axis_tensor)[0];
+      need_resize_out_dims = true;
+    }
+    axis = ComputeAxis(static_cast<int64_t>(axis),
+                       static_cast<int64_t>(ins[0]->dims().size()));
+
+    if (need_resize_out_dims) {
+      const size_t n = ins.size();
+      std::vector<framework::DDim> ins_dims(n);
+      for (size_t i = 0; i < n; i++) {
+        ins_dims[i] = ins[i]->dims();
+      }
+
+      framework::DDim out_dims = ComputeAndCheckShape(true, ins_dims, axis);
+      out->Resize(out_dims);
+    }
     auto place = ctx.GetPlace();
     out->mutable_data<T>(place);
 
@@ -92,10 +147,15 @@ class ConcatGradKernel : public framework::OpKernel<T> {
         }
       }
     }
-    PADDLE_ENFORCE(ins[0], "The input should not be null.");
-    auto axis = ComputeAxis(static_cast<int64_t>(ctx.Attr<int>("axis")),
-                            static_cast<int64_t>(ins[0]->dims().size()));
+    PADDLE_ENFORCE_EQ(ins[0] != nullptr, true, "The input should not be null.");
 
+    auto axis = ctx.Attr<int>("axis");
+    if (ctx.HasInput("AxisTensor")) {
+      auto* axis_tensor = ctx.Input<framework::Tensor>("AxisTensor");
+      axis = GetDataFromTensor<int>(axis_tensor)[0];
+    }
+    axis = ComputeAxis(static_cast<int64_t>(axis),
+                       static_cast<int64_t>(ins[0]->dims().size()));
     // get output tensor that the name is not kEmptyVarName
     std::vector<framework::Tensor*> outputs;
     for (size_t j = 0; j < outs.size(); ++j) {

paddle/fluid/operators/reshape_op.cc

@@ -186,13 +186,16 @@ class ReshapeOp : public framework::OperatorWithKernel {
         output_shape[unk_dim_idx] = -1;
       }
     } else {
-      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);
+      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);
+      }
     }
     return framework::make_ddim(output_shape);
   }

paddle/fluid/operators/split_op.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/split_op.h"
+#include <string>
 
 namespace paddle {
 namespace operators {
@@ -23,8 +24,8 @@ class SplitOp : public framework::OperatorWithKernel {
   using framework::OperatorWithKernel::OperatorWithKernel;
 
   void InferShape(framework::InferShapeContext *ctx) const override {
-    PADDLE_ENFORCE(ctx->HasInput("X"),
-                   "Input(X) of SplitOp should not be null.");
+    PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
+                      "Input(X) of SplitOp should not be null.");
     PADDLE_ENFORCE_GE(ctx->Outputs("Out").size(), 1UL,
                       "Outputs(Out) of SplitOp should not be empty.");
     auto in_dims = ctx->GetInputDim("X");
@@ -34,38 +35,29 @@ class SplitOp : public framework::OperatorWithKernel {
     std::vector<int> sections = static_cast<std::vector<int>>(
         ctx->Attrs().Get<std::vector<int>>("sections"));
     const size_t outs_number = outs_names.size();
-    std::vector<framework::DDim> outs_dims;
-    outs_dims.reserve(outs_number);
-
-    if (num > 0) {
-      int64_t in_axis_dim = in_dims[axis];
-      if (ctx->IsRuntime() || in_axis_dim > 0) {
-        PADDLE_ENFORCE_EQ(in_axis_dim % num, 0,
-                          "tensor split does not result"
-                          " in an equal division");
-        size_t out_axis_dim = in_axis_dim / num;
-        for (size_t i = 0; i < outs_number; ++i) {
-          auto dim = in_dims;
-          dim[axis] = out_axis_dim;
-          outs_dims.push_back(dim);
-        }
-      } else {
-        for (size_t i = 0; i < outs_number; ++i) {
-          auto dim = in_dims;
-          dim[axis] = -1;
-          outs_dims.push_back(dim);
-        }
-      }
-    } else if (sections.size() > 0) {
+
+    if (sections.size() > 0) {
       PADDLE_ENFORCE_EQ(sections.size(), outs_number,
-                        "tensor split sections size"
+                        "tensor split sections size "
                         "should be equal to output size.");
+    }
+
+    if (ctx->HasInput("AxisTensor")) {
+      auto out_dims =
+          framework::make_ddim(std::vector<int>(in_dims.size(), -1));
+      std::vector<framework::DDim> outs_dims(outs_number, out_dims);
+      ctx->SetOutputsDim("Out", outs_dims);
       for (size_t i = 0; i < outs_number; ++i) {
-        auto dim = in_dims;
-        dim[axis] = sections[i];
-        outs_dims.push_back(dim);
+        ctx->ShareLoD("X", "Out", 0, i);
       }
+      return;
     }
+
+    bool each_section_is_known =
+        (sections.size() > 0 && !ctx->HasInputs("SectionsTensorList"));
+
+    auto outs_dims = UpdateOutsDims(ctx->IsRuntime(), each_section_is_known,
+                                    in_dims, num, sections, axis, outs_number);
     ctx->SetOutputsDim("Out", outs_dims);
     if (axis != 0) {
       // Only pass LoD when not spliting along the first dim.
@@ -74,12 +66,41 @@ class SplitOp : public framework::OperatorWithKernel {
       }
     }
   }
+
+ protected:
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext &ctx) const override {
+    return framework::OpKernelType(ctx.Input<framework::LoDTensor>("X")->type(),
+                                   ctx.device_context());
+  }
+
+  framework::OpKernelType GetKernelTypeForVar(
+      const std::string &var_name, const Tensor &tensor,
+      const framework::OpKernelType &expected_kernel_type) const override {
+    if (var_name == "AxisTensor" || var_name == "SectionsTensorList") {
+      return expected_kernel_type;
+    }
+    return framework::OpKernelType(expected_kernel_type.data_type_,
+                                   tensor.place(), tensor.layout());
+  }
 };
 
 class SplitOpMaker : public framework::OpProtoAndCheckerMaker {
  public:
   void Make() override {
     AddInput("X", "(Tensor) Input tensor of the split operator.");
+    AddInput("AxisTensor",
+             "(Tensor) The axis which the input will be splited on. "
+             "It has higher priority than Attr(axis). "
+             "The shape of AxisTensor must be [1]")
+        .AsDispensable();
+    AddInput("SectionsTensorList",
+             "(vector<Tensor<int>>, optional). "
+             "The length of each output along the specified axis. "
+             "It has a higher priority than Attr(sections)."
+             "The shape of the element in vector must be [1].")
+        .AsDuplicable()
+        .AsDispensable();
     AddOutput("Out", "(Tensor) Output tensors of the split operator.")
         .AsDuplicable();
     AddComment(R"DOC(