【Hackathon 6th No.9】fix reshape op with zero-size tensor input in static mode -part (#65639)

NKNaN · web-flow · commit 61bb43a9432e · 2024-07-06T00:48:14.000+08:00
* update reshape op

* fix test

* update reshape_op

* update reshape_op

* update manipulation
diff --git a/paddle/fluid/operators/reshape_op.cc b/paddle/fluid/operators/reshape_op.cc
@@ -139,19 +139,13 @@ class ReshapeOp : public framework::OperatorWithKernel {
                                  const phi::DDim &in_dims) {
     const int64_t in_size = common::product(in_dims);
     auto in_dims_vec = common::vectorize(in_dims);
-    bool all_positive = std::all_of(in_dims_vec.cbegin(),
-                                    in_dims_vec.cend(),
-                                    [](int64_t i) { return i > 0; });
-    // 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) {
+      if (shape[i] == -1) {
+        // only one dimension can be set to -1, whose size will be infered.
         PADDLE_ENFORCE_EQ(
             unk_dim_idx,
             -1,
@@ -161,19 +155,33 @@ class ReshapeOp : public framework::OperatorWithKernel {
                 common::make_ddim(shape),
                 i));
         unk_dim_idx = static_cast<int>(i);
-      } else if (shape[i] == copy_dim_val) {
-        PADDLE_ENFORCE_LT(
-            static_cast<int>(i),
-            in_dims.size(),
-            phi::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.",
-                common::make_ddim(shape),
-                i,
-                in_dims,
-                in_dims.size()));
+        output_shape[i] = shape[i];
+      } else if (shape[i] == 0) {
+        if (in_size == 0) {
+          // zero-sized tensor case
+          // index i could be < in_dims.size(): such as [3, 2, 0] -> [0, 0] is
+          // [0, 0], [3, 2, 0] -> [10, 0] is [10, 0]; index i could be >=
+          // in_dims.size(): such as [3, 2, 0] -> [1, 3, 0, 0] is [1, 3, 0, 0]
+          output_shape[i] = 0;
+        } else {
+          // in other cases 0 means keep in_dims[i] unchanged
+          // index i must be < in_dims.size(): such as [3, 2, 1] -> [0, 0]
+          // is [3, 2] or [3, 2, 1] -> [3, 2, 0] is [3, 2, 1]
+          PADDLE_ENFORCE_LT(
+              static_cast<int>(i),
+              in_dims.size(),
+              phi::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.",
+                  phi::make_ddim(shape),
+                  i,
+                  in_dims,
+                  in_dims.size()));
+          output_shape[i] = in_dims[static_cast<int>(i)];
+        }
+        capacity *= output_shape[i];
       } else {
         PADDLE_ENFORCE_GT(
             shape[i],
@@ -185,25 +193,36 @@ class ReshapeOp : public framework::OperatorWithKernel {
                 common::make_ddim(shape),
                 i,
                 shape[i]));
+        output_shape[i] = shape[i];
+        capacity *= output_shape[i];
       }
+    }
 
-      // NOTE all non-zero values will be converted to True (include negative
-      // value)
-      capacity *= (shape[i] ? shape[i] : in_dims[static_cast<int>(i)]);
-      output_shape[i] = (shape[i] ? static_cast<int64_t>(shape[i])
-                                  : in_dims[static_cast<int>(i)]);
+    if (capacity == 0) {
+      PADDLE_ENFORCE_EQ(in_size,
+                        0,
+                        phi::errors::InvalidArgument(
+                            "Only Zero-Size Tensor'shape can contain 0"));
+      PADDLE_ENFORCE_EQ(unk_dim_idx,
+                        -1,
+                        phi::errors::InvalidArgument(
+                            "can not reshape %s to %s, because the unspecified "
+                            "dimension %i can be any number and is ambiguous",
+                            in_dims,
+                            common::make_ddim(shape),
+                            unk_dim_idx));
     }
 
+    bool no_negative = std::all_of(in_dims_vec.cbegin(),
+                                   in_dims_vec.cend(),
+                                   [](int64_t i) { return i >= 0; });
     if (unk_dim_idx != -1) {
-      if (all_positive) {
-        // 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;
+      // in compile time, no_negative may be False.
+      if (no_negative) {
+        output_shape[unk_dim_idx] = in_size / capacity;
         PADDLE_ENFORCE_EQ(
             output_shape[unk_dim_idx] * capacity,
-            -in_size,
+            in_size,
             phi::errors::InvalidArgument(
                 "The 'shape' attribute in ReshapeOp is invalid. "
                 "The input tensor X'size must be divisible by known "
@@ -215,10 +234,11 @@ class ReshapeOp : public framework::OperatorWithKernel {
                 common::make_ddim(shape),
                 capacity));
       } else {
+        // such as [-1, 8, 3]->[-1, 8], out_shape will remain [-1, 8]
         output_shape[unk_dim_idx] = -1;
       }
     } else {
-      if (all_positive) {
+      if (no_negative) {
         PADDLE_ENFORCE_EQ(
             capacity,
             in_size,
@@ -235,24 +255,6 @@ class ReshapeOp : public framework::OperatorWithKernel {
       }
     }
 
-    // support reshape with zero-input(input tensor with product(shape) == 0)
-    // by now we require that if the input tensor is zero shape, the target
-    // shape of output must be zero
-    if (in_size == 0) {
-      PADDLE_ENFORCE_LE(
-          capacity,
-          in_size,
-          phi::errors::InvalidArgument(
-              "The 'shape' in ReshapeOp is invalid. "
-              "The input tensor X's shape = [%s], X's capacity = %d."
-              "But the target shape of Out is [%s],  the "
-              "capacity of 'Out' is %d.",
-              in_dims,
-              in_size,
-              common::make_ddim(shape),
-              capacity));
-    }
-
     return common::make_ddim(output_shape);
   }
 
diff --git a/paddle/phi/kernels/reshape_kernel.cc b/paddle/phi/kernels/reshape_kernel.cc
@@ -32,10 +32,7 @@ void ReshapeInferKernel(const Context& dev_ctx,
                         DenseTensor* out) {
   MetaTensor meta_out(out);
   InferMetaFromVecValue(x, shape.GetData(), &meta_out);
-  // Zero-Size Tensor
-  if (x.numel() == 0) {
-    return;
-  }
+
   if (x.initialized() && x.Holder() == out->Holder()) {
     dev_ctx.Alloc(out, x.dtype());
     return;
@@ -57,9 +54,7 @@ void ReshapeInferKernel<phi::XPUContext>(const XPUContext& dev_ctx,
                                          DenseTensor* out) {
   MetaTensor meta_out(out);
   InferMetaFromVecValue(x, shape.GetData(), &meta_out);
-  if (x.numel() == 0) {
-    return;
-  }
+
   if (x.initialized() && x.Holder() == out->Holder()) {
     dev_ctx.Alloc(out, x.dtype());
     return;
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import math
 from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence, overload
 
 import numpy as np
@@ -4721,12 +4722,13 @@ def get_attr_shape(list_shape):
                     )
                     unk_dim_idx = dim_idx
                 elif dim_size == 0:
-                    assert dim_idx < len(x.shape), (
-                        "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."
-                        % (dim_idx, len(x.shape))
-                    )
+                    if math.prod(x.shape):
+                        assert dim_idx < len(x.shape), (
+                            "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."
+                            % (dim_idx, len(x.shape))
+                        )
                 else:
                     assert dim_size > 0, (
                         "Each dimension value of 'shape' in reshape must not "
diff --git a/test/legacy_test/test_zero_size_tensor.py b/test/legacy_test/test_zero_size_tensor.py
@@ -19,7 +19,10 @@
 
 import unittest
 
+import numpy as np
+
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 
 # Use to test zero-size of Sundry API, which is unique and can not be classified
@@ -39,7 +42,7 @@ def test_numpy(self):
         self.assertEqual(out.shape, (0, 2))
         self.assertEqual(out.size, 0)
 
-    def test_reshape(self):
+    def test_reshape_dygraph(self):
         # case 1
         x1 = paddle.rand([0, 2])
         x1.stop_gradient = False
@@ -77,6 +80,40 @@ def test_reshape(self):
         with self.assertRaises(ValueError):
             out4 = paddle.reshape(x5, [2, 0, -1])
 
+    @test_with_pir_api
+    def test_reshape_static(self):
+        paddle.enable_static()
+        place = paddle.CPUPlace()
+        if paddle.is_compiled_with_cuda():
+            place = paddle.CUDAPlace(0)
+
+        input_cases = [
+            # (x, new_shape, desired_shape)
+            (np.random.rand(0, 2), [-1], [0]),
+            (np.random.rand(0, 2), [2, -1], [2, 0]),
+            (np.random.rand(0, 2), [2, 3, 0], [2, 3, 0]),
+            (np.random.rand(0, 2), [0], [0]),
+        ]
+        for case in input_cases:
+            data_np, new_shape, desired_shape = case
+            startup_program = paddle.static.Program()
+            main_program = paddle.static.Program()
+            executor = paddle.static.Executor(place)
+            with paddle.static.program_guard(main_program, startup_program):
+                x = paddle.static.data(
+                    name="x", shape=data_np.shape, dtype='float64'
+                )
+                out = paddle.reshape(x, new_shape)
+                fetch_list = [out]
+                feeds = {'x': data_np}
+
+                executor.run(startup_program)
+                res = executor.run(
+                    main_program, feed=feeds, fetch_list=fetch_list
+                )
+            np.testing.assert_equal(res[0], np.random.rand(*desired_shape))
+        paddle.disable_static()
+
 
 if __name__ == "__main__":
     unittest.main()
