Fix 0-size tensor support in mean when out->numel() > 0, add test.

Author: cangtianhuang

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;

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>(

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,

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,

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 = "comp"
+
+    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()