[0-size Tensor Job2 No.78、98] Add 0-size Tensor support for sum (#73355)

co63oc · web-flow · commit d0df4b205524 · 2025-06-18T15:12:20.000+08:00
* Fix

* Fix

* Fix
diff --git a/paddle/phi/kernels/cpu/reduce_sum_kernel.cc b/paddle/phi/kernels/cpu/reduce_sum_kernel.cc
@@ -37,11 +37,23 @@ void SumRawKernel(const Context& dev_ctx,
   }
   if (x.numel() == 0) {
     dev_ctx.template Alloc<T>(out);
-    FullKernel<T, Context>(dev_ctx,
-                           phi::IntArray(common::vectorize(out->dims())),
-                           0,
-                           out_dtype,
-                           out);
+    // When out_dtype is DataType::UNDEFINED and input is int32 or bool,
+    // result is int64, but FullKernel out_dtype parameter is not used, we need
+    // to set int64 explicitly.
+    if (out_dtype == DataType::INT64) {
+      FullKernel<int64_t, Context>(
+          dev_ctx,
+          phi::IntArray(common::vectorize(out->dims())),
+          0,
+          out_dtype,  // not used
+          out);
+    } else {
+      FullKernel<T, Context>(dev_ctx,
+                             phi::IntArray(common::vectorize(out->dims())),
+                             0,
+                             out_dtype,  // not used
+                             out);
+    }
     return;
   }
   if constexpr (std::is_same_v<T, phi::dtype::float16> ||
diff --git a/paddle/phi/kernels/kps/reduce_kernel.cu b/paddle/phi/kernels/kps/reduce_kernel.cu
@@ -212,11 +212,20 @@ void SumRawKernel(const Context& dev_ctx,
   }
   if (x.numel() == 0) {
     dev_ctx.template Alloc<T>(out);
-    FullKernel<T, Context>(dev_ctx,
-                           phi::IntArray(common::vectorize(out->dims())),
-                           0,
-                           out_dtype,
-                           out);
+    if (out_dtype == DataType::INT64) {
+      FullKernel<int64_t, Context>(
+          dev_ctx,
+          phi::IntArray(common::vectorize(out->dims())),
+          0,
+          out_dtype,  // not used
+          out);
+    } else {
+      FullKernel<T, Context>(dev_ctx,
+                             phi::IntArray(common::vectorize(out->dims())),
+                             0,
+                             out_dtype,  // not used
+                             out);
+    }
     return;
   }
 
diff --git a/paddle/phi/kernels/xpu/reduce_sum_kernel.cc b/paddle/phi/kernels/xpu/reduce_sum_kernel.cc
@@ -34,11 +34,20 @@ void SumRawKernel(const Context& dev_ctx,
   }
   if (x.numel() == 0) {
     dev_ctx.template Alloc<T>(out);
-    FullKernel<T, Context>(dev_ctx,
-                           phi::IntArray(common::vectorize(out->dims())),
-                           0,
-                           out_dtype,
-                           out);
+    if (out_dtype == DataType::INT64) {
+      FullKernel<int64_t, Context>(
+          dev_ctx,
+          phi::IntArray(common::vectorize(out->dims())),
+          0,
+          out_dtype,  // not used
+          out);
+    } else {
+      FullKernel<T, Context>(dev_ctx,
+                             phi::IntArray(common::vectorize(out->dims())),
+                             0,
+                             out_dtype,  // not used
+                             out);
+    }
     return;
   }
   XPUReduce<Context, T, phi::SumFunctor>(
diff --git a/test/legacy_test/test_sum_op.py b/test/legacy_test/test_sum_op.py
@@ -923,6 +923,48 @@ def test_warnings(self):
                 os.environ["FLAGS_print_extra_attrs"] = '0'
 
 
+class TestSum_BoolToInt64_ZeroSize(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(123)
+        self.shape = [3, 0, 2]
+        self.places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            self.places.append(base.CUDAPlace(0))
+
+    def check_result(
+        self, dygraph_result, expected_result, axis, keepdim, dtype, place
+    ):
+        self.assertTrue(
+            (dygraph_result == expected_result).all(),
+            f"Shape: {self.shape}, Axis: {axis}, Keepdim: {keepdim}, Dtype: {dtype}, Place: {place}",
+        )
+
+    def _test_dygraph(self, place, axis, keepdim, dtype):
+        with dygraph_guard():
+            x_np = np.random.random(self.shape).astype(dtype)
+            x = paddle.to_tensor(x_np)
+            x.stop_gradient = False
+            dygraph_result = paddle.sum(x, axis=axis, keepdim=keepdim)
+            expected_result = np.sum(x_np, axis=axis, keepdims=keepdim)
+            self.check_result(
+                dygraph_result.numpy(),
+                expected_result,
+                axis,
+                keepdim,
+                dtype,
+                place,
+            )
+            paddle.sum(dygraph_result).backward()
+            np.testing.assert_allclose(x.grad.shape, x.shape)
+
+    def test_zero_size(self):
+        keepdims_options = [True, False]
+        for place in self.places:
+            for keepdim in keepdims_options:
+                self._test_dygraph(place, None, keepdim, "bool")
+                self._test_dygraph(place, None, keepdim, "int32")
+
+
 if __name__ == "__main__":
     enable_static()
     unittest.main()
