[0-size Tensor Job2 No.46] Add 0-size Tensor support for masked_scatter (#73384)

co63oc · web-flow · commit 2f56bf416a3e · 2025-06-20T16:50:43.000+08:00
diff --git a/paddle/phi/kernels/xpu/clip_grad_kernel.cc b/paddle/phi/kernels/xpu/clip_grad_kernel.cc
@@ -27,6 +27,7 @@ void ClipGradKernel(const Context& dev_ctx,
                     const Scalar& max,
                     DenseTensor* x_grad) {
   dev_ctx.template Alloc<T>(x_grad);
+  if (x_grad && x_grad->numel() == 0) return;
   using XPUDataType = typename XPUTypeTrait<T>::Type;
   int r =
       xpu::clamp_grad(dev_ctx.x_context(),
diff --git a/paddle/phi/kernels/xpu/clip_kernel.cc b/paddle/phi/kernels/xpu/clip_kernel.cc
@@ -41,6 +41,7 @@ void ClipKernel(const Context& dev_ctx,
                               static_cast<float>(max_)));
 
   dev_ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) return;
   using XPUDataType = typename XPUTypeTrait<T>::Type;
   auto x_data = reinterpret_cast<const XPUDataType*>(x.data<T>());
   auto out_data = reinterpret_cast<XPUDataType*>(out->data<T>());
diff --git a/paddle/phi/kernels/xpu/logical_kernel.cc b/paddle/phi/kernels/xpu/logical_kernel.cc
@@ -24,6 +24,7 @@ void LogicalNotKernel(const Context& dev_ctx,
                       const DenseTensor& x,
                       DenseTensor* out) {
   dev_ctx.template Alloc<bool>(out);
+  if (out && out->numel() == 0) return;
   int r = xpu::logical_not(
       dev_ctx.x_context(), x.data<T>(), out->data<T>(), x.numel());
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "logical_not");
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -5158,7 +5158,7 @@ def masked_scatter(
     zeros_like_x = paddle.zeros_like(x, dtype=int)
     mask = paddle.add(paddle.cast(mask, dtype="int"), zeros_like_x)
     mask_prefix = paddle.clip(mask.cumsum() - 1, min=0)
-    if in_dynamic_mode():
+    if in_dynamic_mode() and mask_prefix.numel() != 0:
         assert (
             mask_prefix[-1] <= value.numel()
         ), f'mask true nums must be <= value size, but got mask true nums is {mask_prefix[-1].item()}, value size is {value.numel().item()}'
diff --git a/test/legacy_test/test_masked_scatter.py b/test/legacy_test/test_masked_scatter.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import unittest
 
 import numpy as np
@@ -65,6 +66,7 @@ def test_dtype_error(self):
             paddle.masked_scatter(x, mask, value)
 
     def test_numel_error(self):
+        paddle.disable_static()
         self.value_np = np.random.randn(5, 5).astype(self.dtype)
         x = paddle.to_tensor(self.x_np, dtype=self.dtype)
         mask = paddle.to_tensor(self.mask_np).astype('bool')
@@ -309,6 +311,52 @@ def init(self):
         self.value_shape = (300, 300)
 
 
+class TestMaskedScatterAPI_ZeroSize(unittest.TestCase):
+    def setUp(self):
+        self.init()
+
+        self.x_np = np.random.random(self.x_shape).astype(self.dtype)
+        self.mask_np = np.array(
+            np.random.randint(2, size=self.mask_shape), dtype="bool"
+        )
+
+        self.value_np = np.random.randn(*self.value_shape).astype(self.dtype)
+        self.out_np = np_masked_scatter(self.x_np, self.mask_np, self.value_np)
+
+        places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not core.is_compiled_with_cuda()
+        ):
+            places.append(base.CPUPlace())
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        self.places = places
+
+    def init(self):
+        self.x_shape = (3, 0)
+        self.mask_shape = self.x_shape
+        self.dtype = "float32"
+        self.value_shape = (300, 300)
+
+    def _test_dygraph(self, place):
+        paddle.disable_static(place)
+        x = paddle.to_tensor(self.x_np, dtype=self.dtype)
+        x.stop_gradient = False
+        mask = paddle.to_tensor(self.mask_np).astype('bool')
+        value = paddle.to_tensor(self.value_np, dtype=self.dtype)
+        result = paddle.masked_scatter(x, mask, value)
+        np.testing.assert_allclose(self.out_np, result.numpy(), rtol=1e-05)
+        paddle.sum(result).backward()
+        np.testing.assert_allclose(x.grad.shape, x.shape)
+        paddle.enable_static()
+
+    def test_dygraph(self):
+        for place in self.places:
+            self._test_dygraph(place)
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()
