[0-size Tensor No.43、131] Add 0-size Tensor support for paddle.diag/masked_fill (#73570)

co63oc · web-flow · commit e991019e58a3 · 2025-07-01T17:13:04.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -899,6 +899,9 @@ bool DiagOpInferSymbolicShape(pir::Operation *op,
           size_ = x_shape[1].dyn_cast<int64_t>();
         }
       }
+      if (size_ < 0) {
+        size_ = 0;
+      }
       infer_context->SetShapeOrDataForValue(
           op->result(0), symbol::TensorShapeOrDataDimExprs({size_}));
     } else {
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -955,6 +955,9 @@ void DiagInferMeta(const MetaTensor& x,
         size_ = x_dims[1];
       }
     }
+    if (size_ < 0) {
+      size_ = 0;
+    }
     out->set_dims({size_});
     out->set_dtype(x.dtype());
   } else {
diff --git a/paddle/phi/kernels/cpu/diag_grad_kernel.cc b/paddle/phi/kernels/cpu/diag_grad_kernel.cc
@@ -28,6 +28,7 @@ void DiagGradKernel(const Context& dev_ctx,
                     int offset,
                     DenseTensor* x_grad) {
   T* dx_data = dev_ctx.template Alloc<T>(x_grad);
+  if (x_grad && x_grad->numel() == 0) return;
   const T* dout_data = out_grad.data<T>();
   auto dx_dims = x_grad->dims();
   auto dout_dims = out_grad.dims();
diff --git a/paddle/phi/kernels/cpu/diag_kernel.cc b/paddle/phi/kernels/cpu/diag_kernel.cc
@@ -30,6 +30,7 @@ void DiagKernel(const Context& dev_ctx,
   auto* x_data = x.data<T>();
   auto x_dims = x.dims();
   T* out_data = dev_ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) return;
   auto out_dims = out->dims();
 
   int64_t i = 0;
diff --git a/paddle/phi/kernels/cpu/masked_fill_grad_kernel.cc b/paddle/phi/kernels/cpu/masked_fill_grad_kernel.cc
@@ -20,9 +20,9 @@
 #include "paddle/phi/kernels/empty_kernel.h"
 #include "paddle/phi/kernels/expand_grad_kernel.h"
 #include "paddle/phi/kernels/expand_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/common_infer_shape_functions.h"
 #include "paddle/phi/kernels/funcs/common_shape.h"
-
 namespace phi {
 
 template <typename T, typename Context>
@@ -33,6 +33,19 @@ void MaskedFillGradKernel(const Context& dev_ctx,
                           const DenseTensor& out_grad,
                           DenseTensor* x_grad,
                           DenseTensor* v_grad) {
+  if (out_grad.numel() == 0 || mask.numel() == 0) {
+    // x shape [2, 1, 3], mask shape [2, 0, 3], x_grad shape [2, 1, 3]
+    if (x_grad) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+    }
+    if (v_grad) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(v_grad->dims())), 0, v_grad);
+    }
+    return;
+  }
+
   auto x_grad_dims = x_grad->dims();
   auto mask_dims = mask.dims();
   bool expand_x = false;
diff --git a/paddle/phi/kernels/cpu/masked_fill_kernel.cc b/paddle/phi/kernels/cpu/masked_fill_kernel.cc
@@ -29,6 +29,11 @@ void MaskedFillKernel(const Context& dev_ctx,
                       const DenseTensor& mask,
                       const DenseTensor& value,
                       DenseTensor* out) {
+  if (x.numel() == 0 || mask.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
+
   auto x_dims = x.dims();
   auto mask_dims = mask.dims();
 
diff --git a/paddle/phi/kernels/gpu/diag_grad_kernel.cu b/paddle/phi/kernels/gpu/diag_grad_kernel.cu
@@ -57,6 +57,7 @@ void DiagGradKernel(const Context& dev_ctx,
                     int offset,
                     DenseTensor* x_grad) {
   T* dx_data = dev_ctx.template Alloc<T>(x_grad);
+  if (x_grad && x_grad->numel() == 0) return;
   auto* dout_data = out_grad.data<T>();
   auto dx_dims = x_grad->dims();
   auto dout_dims = out_grad.dims();
diff --git a/paddle/phi/kernels/gpu/diag_kernel.cu b/paddle/phi/kernels/gpu/diag_kernel.cu
@@ -63,6 +63,7 @@ void DiagKernel(const Context& dev_ctx,
   auto* x_data = x.data<T>();
   auto x_dims = x.dims();
   T* out_data = dev_ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) return;
   auto out_dims = out->dims();
 
   auto GetBlockGridSize = [&dev_ctx](int64_t size) {
diff --git a/paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu b/paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu
@@ -307,13 +307,14 @@ void MaskedFillGradKernel(const Context& dev_ctx,
                           DenseTensor* x_grad,
                           DenseTensor* v_grad) {
   if (out_grad.numel() == 0 || mask.numel() == 0) {
-    if (x_grad != nullptr) {
-      x_grad->Resize({0});
-      dev_ctx.template Alloc<T>(x_grad);
+    // x shape [2, 1, 3], mask shape [2, 0, 3], x_grad shape [2, 1, 3]
+    if (x_grad) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
     }
-    if (v_grad != nullptr) {
-      v_grad->Resize({0});
-      dev_ctx.template Alloc<T>(v_grad);
+    if (v_grad) {
+      phi::Full<T, Context>(
+          dev_ctx, phi::IntArray(common::vectorize(v_grad->dims())), 0, v_grad);
     }
     return;
   }
diff --git a/paddle/phi/kernels/gpu/masked_fill_kernel.cu b/paddle/phi/kernels/gpu/masked_fill_kernel.cu
@@ -211,7 +211,6 @@ void MaskedFillKernel(const Context& dev_ctx,
                       const DenseTensor& value,
                       DenseTensor* out) {
   if (x.numel() == 0 || mask.numel() == 0) {
-    out->Resize({0});
     dev_ctx.template Alloc<T>(out);
     return;
   }
diff --git a/paddle/phi/kernels/xpu/diag_kernel.cc b/paddle/phi/kernels/xpu/diag_kernel.cc
@@ -29,6 +29,7 @@ void DiagKernel(const Context& dev_ctx,
   using XPUType = typename XPUTypeTrait<T>::Type;
   auto* x_data = reinterpret_cast<const XPUType*>(x.data<T>());
   dev_ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) return;
   auto* out_data = reinterpret_cast<XPUType*>(out->data<T>());
 
   auto x_shape = common::vectorize<int64_t>(x.dims());
diff --git a/test/legacy_test/test_diag_v2.py b/test/legacy_test/test_diag_v2.py
@@ -400,6 +400,44 @@ def init_config(self):
         self.out = np.diag(self.x, self.offset)
 
 
+class TestDiagV2Op_ZeroSize(OpTest):
+    def setUp(self):
+        self.op_type = "diag_v2"
+        self.python_api = paddle.diag
+        self.public_python_api = paddle.diag
+
+        self.init_dtype()
+        self.init_attrs()
+        self.init_input_output()
+        self.set_input_output()
+
+    def init_dtype(self):
+        self.dtype = np.float64
+
+    def init_attrs(self):
+        self.offset = 1
+        self.padding_value = 0.0
+
+    def init_input_output(self):
+        self.x = np.random.rand(10, 0).astype(self.dtype)
+        self.out = np.diag(self.x, self.offset)
+
+    def set_input_output(self):
+        self.attrs = {
+            'offset': self.offset,
+            'padding_value': self.padding_value,
+        }
+
+        self.inputs = {'X': self.x}
+        self.outputs = {'Out': self.out}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True)
+
+    def test_check_grad(self):
+        self.check_grad(['X'], 'Out', check_pir=True)
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()
diff --git a/test/legacy_test/test_masked_fill.py b/test/legacy_test/test_masked_fill.py
@@ -145,6 +145,7 @@ def setUp(self):
         self.dtype = "float32"
 
     def test_backward(self):
+        paddle.disable_static()
         expected_np = np.array(
             [[2, 1, 1], [2, 1, 1], [2, 1, 1], [2, 1, 1]]
         ).astype('float32')
@@ -358,6 +359,50 @@ def init(self):
         self.scalar_value = False
 
 
+class TestMaskedFillAPI_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(1).astype(self.dtype)
+        self.out_np = np_masked_fill(self.x_np, self.mask_np, self.value_np)
+
+    def init(self):
+        self.x_shape = (0, 3)
+        self.mask_shape = self.x_shape
+        self.dtype = "float32"
+        self.scalar_value = False
+
+    def test_dygraph(self):
+        paddle.disable_static()
+        x = paddle.to_tensor(self.x_np, dtype=self.dtype)
+        x.stop_gradient = False
+        mask = paddle.to_tensor(self.mask_np).astype('bool')
+        if self.scalar_value:
+            value = self.value_np[0]
+        else:
+            value = paddle.to_tensor(self.value_np, dtype=self.dtype)
+        result = paddle.masked_fill(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)
+        np.testing.assert_allclose(x.grad.numpy(), np.zeros(x.shape))
+
+
+class TestMaskedFillAPI_ZeroSize2(TestMaskedFillAPI_ZeroSize):
+    # x_grad shape [2, 3], filled with 0.
+    def init(self):
+        self.x_shape = (1, 3)
+        self.mask_shape = (0, 3)
+        self.dtype = "float32"
+        self.scalar_value = False
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -899,6 +899,9 @@ bool DiagOpInferSymbolicShape(pir::Operation *op,`
`899`	`899`	`size_ = x_shape[1].dyn_cast<int64_t>();`
`900`	`900`	`}`
`901`	`901`	`}`
	`902`	`+ if (size_ < 0) {`
	`903`	`+ size_ = 0;`
	`904`	`+ }`
`902`	`905`	`infer_context->SetShapeOrDataForValue(`
`903`	`906`	`op->result(0), symbol::TensorShapeOrDataDimExprs({size_}));`
`904`	`907`	`} else {`
Original file line number	Diff line number	Diff line change
`@@ -955,6 +955,9 @@ void DiagInferMeta(const MetaTensor& x,`
`955`	`955`	`size_ = x_dims[1];`
`956`	`956`	`}`
`957`	`957`	`}`
	`958`	`+ if (size_ < 0) {`
	`959`	`+ size_ = 0;`
	`960`	`+ }`
`958`	`961`	`out->set_dims({size_});`
`959`	`962`	`out->set_dtype(x.dtype());`
`960`	`963`	`} else {`
Original file line number	Diff line number	Diff line change
`@@ -211,7 +211,6 @@ void MaskedFillKernel(const Context& dev_ctx,`
`211`	`211`	`const DenseTensor& value,`
`212`	`212`	`DenseTensor* out) {`
`213`	`213`	`if (x.numel() == 0 \|\| mask.numel() == 0) {`
`214`		`- out->Resize({0});`
`215`	`214`	`dev_ctx.template Alloc<T>(out);`
`216`	`215`	`return;`
`217`	`216`	`}`