[0-Size Tensor问题修复] paddle.mean (#71961)

cangtianhuang · web-flow · commit fd377187aedf · 2025-04-01T15:34:04.000+08:00
* Fix 0-size tensor support in mean when out-&gt;numel() &gt; 0, add test.

* Restore GT check in reduce_function.

* Fix bug.
diff --git a/paddle/phi/kernels/cpu/reduce_mean_kernel.cc b/paddle/phi/kernels/cpu/reduce_mean_kernel.cc
@@ -29,12 +29,7 @@ void MeanRawKernel(const Context& dev_ctx,
                    bool keep_dim,
                    bool reduce_all,
                    DenseTensor* out) {
-  if (out && out->numel() == 0) {
-    dev_ctx.template Alloc<T>(out);
-    return;
-  }
-
-  if (x.numel() == 0 && out && out->dims().size() == 0) {
+  if (x.numel() == 0) {
     phi::Full<T, Context>(
         dev_ctx, phi::IntArray(common::vectorize(out->dims())), NAN, out);
     return;
diff --git a/paddle/phi/kernels/funcs/reduce_function.h b/paddle/phi/kernels/funcs/reduce_function.h
@@ -1357,11 +1357,10 @@ void ReduceKernelImpl(const Context& dev_ctx,
                       const std::vector<int64_t>& dims,
                       bool keep_dim,
                       bool reduce_all) {
-  PADDLE_ENFORCE_GE(input.numel(),
+  PADDLE_ENFORCE_GT(input.numel(),
                     0,
                     common::errors::InvalidArgument(
-                        "The input size (numel) of reduce op should be larger "
-                        "than or equal to 0"));
+                        "Tensor need be reduced must not empty."));
 
   dev_ctx.template Alloc<OutT>(output);
 
diff --git a/paddle/phi/kernels/kps/reduce_kernel.cu b/paddle/phi/kernels/kps/reduce_kernel.cu
@@ -119,16 +119,12 @@ void MeanRawKernel(const Context& dev_ctx,
                    bool keep_dim,
                    bool reduce_all,
                    DenseTensor* out) {
-  if (out && out->numel() == 0) {
-    dev_ctx.template Alloc<T>(out);
-    return;
-  }
-
-  if (x.numel() == 0 && out && out->dims().size() == 0) {
+  if (x.numel() == 0) {
     phi::Full<T, Context>(
         dev_ctx, phi::IntArray(common::vectorize(out->dims())), NAN, out);
     return;
   }
+
   reduce_all = recompute_reduce_all(x, dims, reduce_all);
   auto out_dtype = x.dtype();
   phi::Reduce<T, kps::AddFunctor, kps::IdentityFunctor, true>(
diff --git a/paddle/phi/kernels/onednn/reduce_mean_kernel.cc b/paddle/phi/kernels/onednn/reduce_mean_kernel.cc
@@ -25,16 +25,12 @@ void MeanRawKernel(const Context& dev_ctx,
                    bool keep_dim,
                    bool reduce_all,
                    DenseTensor* out) {
-  if (out && out->numel() == 0) {
-    dev_ctx.template Alloc<T>(out);
-    return;
-  }
-
-  if (x.numel() == 0 && out && out->dims().size() == 0) {
+  if (x.numel() == 0) {
     phi::Full<T, Context>(
         dev_ctx, phi::IntArray(common::vectorize(out->dims())), NAN, out);
     return;
   }
+
   reduce_all = recompute_reduce_all(x, dims, reduce_all);
   ReduceKernel<T, Context>(dev_ctx,
                            x,
diff --git a/paddle/phi/kernels/xpu/reduce_mean_kernel.cc b/paddle/phi/kernels/xpu/reduce_mean_kernel.cc
@@ -29,16 +29,12 @@ void MeanRawKernel(const Context& dev_ctx,
                    bool keep_dim,
                    bool reduce_all,
                    DenseTensor* out) {
-  if (out && out->numel() == 0) {
-    dev_ctx.template Alloc<T>(out);
-    return;
-  }
-
-  if (x.numel() == 0 && out && out->dims().size() == 0) {
+  if (x.numel() == 0) {
     phi::Full<T, Context>(
         dev_ctx, phi::IntArray(common::vectorize(out->dims())), NAN, out);
     return;
   }
+
   reduce_all = recompute_reduce_all(x, dims, reduce_all);
   using XPUType = typename XPUTypeTrait<T>::Type;
   auto f = [](xpu::Context* ctx,
diff --git a/test/legacy_test/test_mean_op.py b/test/legacy_test/test_mean_op.py
@@ -1000,6 +1000,37 @@ def test_check_grad(self):
             )
 
 
+class TestMeanOp_ZeroSize3(OpTest):
+    def setUp(self):
+        self.op_type = 'mean'
+        self.python_api = paddle.mean
+        self.init_prim_type()
+        self.dtype = 'float64'
+        self.shape = [2, 0, 4]
+        self.axis = 1
+        self.keepdim = False
+        self.set_attrs()
+
+        self.inputs = {'X': np.array([], dtype=self.dtype).reshape(self.shape)}
+        self.outputs = {
+            'Out': np.mean(
+                self.inputs["X"], axis=self.axis, keepdims=self.keepdim
+            )
+        }
+
+    def set_attrs(self):
+        pass
+
+    def init_prim_type(self):
+        self.prim_op_type = "prim"
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, equal_nan=True)
+
+    def test_check_grad(self):
+        self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()
