[0-size Tensor Job2 No.87] Add 0-size Tensor support for masked_fill

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 {

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 {

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

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;

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;

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

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

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) {

paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu

@@ -303,13 +303,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;
   }

paddle/phi/kernels/gpu/masked_fill_kernel.cu

@@ -209,7 +209,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;
   }

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

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

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