【Paddle Tensor 第二期】使paddle.all支持0-size输入 (#70228)

aquagull · web-flow · commit 02b8e3a569f1 · 2024-12-16T19:39:02.000+08:00
* update

* update

* fix

* add test

* delete duplicate code

* fix
diff --git a/paddle/phi/kernels/reduce_all_kernel.cc b/paddle/phi/kernels/reduce_all_kernel.cc
@@ -13,10 +13,14 @@
 // limitations under the License.
 
 #include "paddle/phi/kernels/reduce_all_kernel.h"
-
+#include "glog/logging.h"
 #include "paddle/phi/backends/all_context.h"
 #include "paddle/phi/common/complex.h"
+#include "paddle/phi/common/int_array.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/meta_tensor.h"
+#include "paddle/phi/infermeta/unary.h"
+#include "paddle/phi/kernels/full_kernel.h"
 
 using complex64 = ::phi::dtype::complex<float>;
 using complex128 = ::phi::dtype::complex<double>;
@@ -31,10 +35,19 @@ void AllKernel(const Context& dev_ctx,
                DenseTensor* out) {
   auto x_dim = x.dims();
   for (int i = 0; i < x_dim.size(); i++) {
-    PADDLE_ENFORCE_LT(0,
-                      x_dim[i],
-                      errors::InvalidArgument(
-                          "The dims of Input(X) should be greater than 0."));
+    PADDLE_ENFORCE_LE(
+        0,
+        x_dim[i],
+        errors::InvalidArgument(
+            "The dims of Input(X) should be greater than or equal to 0."));
+  }
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<bool>(out);
+    if (out->numel() > 0) {
+      std::vector<int64_t> vec_dims = common::vectorize(out->dims());
+      phi::Full<bool, Context>(dev_ctx, phi::IntArray(vec_dims), 1, out);
+    }
+    return;
   }
   bool reduce_all = recompute_reduce_all(x, dims);
   AllRawKernel<T>(dev_ctx, x, dims, keep_dim, reduce_all, out);
diff --git a/test/legacy_test/test_reduce_op.py b/test/legacy_test/test_reduce_op.py
@@ -17,7 +17,7 @@
 
 import numpy as np
 from op_test import OpTest, convert_float_to_uint16, skip_check_grad_ci
-from utils import static_guard
+from utils import dygraph_guard, static_guard
 
 import paddle
 from paddle import base
@@ -2247,16 +2247,93 @@ def test_dygraph(self):
         paddle.enable_static()
 
 
-class TestAllZeroError(unittest.TestCase):
-    def test_errors(self):
-        with paddle.base.dygraph.guard():
+class TestAllZero(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(123)
+        self.shape = [1, 0, 2]
+        self.dtypes = [
+            "bool",
+            "float32",
+            "float64",
+            "int32",
+            "complex64",
+            "complex128",
+        ]
+        self.places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            self.places.append(base.CUDAPlace(0))
+
+    def calculate_expected_result(self, x_np, axis, keepdim):
+        expected_result = np.all(x_np, axis=axis, keepdims=keepdim)
+        return expected_result
+
+    def check_result(
+        self, static_result, expected_result, axis, keepdim, dtype, place
+    ):
+        self.assertTrue(
+            (static_result == expected_result).all(),
+            f"Static Mode - Shape: {self.shape}, Axis: {axis}, Keepdim: {keepdim}, Dtype: {dtype}, Place: {place}",
+        )
+
+    def _test_static(self, place, axis, keepdim, dtype):
+        with static_guard():
+            with base.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                input = paddle.static.data(
+                    name="x", shape=self.shape, dtype=dtype
+                )
+                result = paddle.all(x=input, axis=axis, keepdim=keepdim)
+                x_np = np.zeros(self.shape, dtype=dtype)
 
-            def test_0_size():
-                array = np.array([], dtype=np.float32)
-                x = paddle.to_tensor(np.reshape(array, [0, 0, 0]), dtype='bool')
-                paddle.all(x, axis=1)
+                exe = base.Executor(place)
+                fetches = exe.run(
+                    feed={"x": x_np},
+                    fetch_list=[result],
+                )
+                expected_result = self.calculate_expected_result(
+                    x_np, axis, keepdim
+                )
+                self.check_result(
+                    fetches[0], expected_result, axis, keepdim, dtype, place
+                )
 
-            self.assertRaises(ValueError, test_0_size)
+    def _test_dygraph(self, place, axis, keepdim, dtype):
+        with dygraph_guard():
+            x_np = np.zeros(self.shape, dtype=dtype)
+            x = paddle.to_tensor(x_np)
+            dygraph_result = paddle.all(x, axis=axis, keepdim=keepdim).numpy()
+            expected_result = self.calculate_expected_result(
+                x_np, axis, keepdim
+            )
+            self.check_result(
+                dygraph_result, expected_result, axis, keepdim, dtype, place
+            )
+
+    def _test_all(self, place, axis, keepdim, dtype):
+        self._test_dygraph(place, axis, keepdim, dtype)
+        self._test_static(place, axis, keepdim, dtype)
+
+    def test_zero_size(self):
+        axes_options = [
+            None,
+            0,
+            1,
+            2,
+            -1,
+            -2,
+            (),
+            (0, 1),
+            (0, 2),
+            (1, 2),
+            (-1, -2),
+        ]
+        keepdims_options = [True, False]
+        for place in self.places:
+            for dtype in self.dtypes:
+                for axis in axes_options:
+                    for keepdim in keepdims_options:
+                        self._test_all(place, axis, keepdim, dtype)
 
 
 if __name__ == '__main__':
