[0-size Tensor No.303] Add 0-size Tensor support for fill_diagonal_

paddle/phi/kernels/cpu/fill_diagonal_grad_kernel.cc

@@ -28,6 +28,7 @@ void FillDiagonalGradKernel(const Context& ctx,
                             DenseTensor* x_grad) {
   if (x_grad) {
     T* data = ctx.template Alloc<T>(x_grad);
+    if (x_grad->numel() == 0) return;
     phi::Copy(ctx, out_grad, ctx.GetPlace(), false, x_grad);
 
     auto dx_dims = x_grad->dims();

paddle/phi/kernels/cpu/fill_diagonal_kernel.cc

@@ -30,6 +30,8 @@ void FillDiagonalKernel(const Context& ctx,
   T temp_var = static_cast<T>(value);
 
   T* out_data = ctx.template Alloc<T>(out);
+  if (out && out->numel() == 0) return;
+
   phi::Copy(ctx, x, ctx.GetPlace(), false, out);
 
   auto out_dims = out->dims();

paddle/phi/kernels/gpu/fill_diagonal_grad_kernel.cu

@@ -52,6 +52,7 @@ void FillDiagonalGradKernel(const Context& ctx,
                             DenseTensor* x_grad) {
   const int64_t kMaxBlockDim = 512;
   auto* in_data = ctx.template Alloc<T>(x_grad);
+  if (x_grad && x_grad->numel() == 0) return;
 
   phi::Copy(ctx, out_grad, ctx.GetPlace(), false, x_grad);
 

paddle/phi/kernels/gpu/fill_diagonal_kernel.cu

@@ -50,6 +50,10 @@ void FillDiagonalKernel(const Context& ctx,
                         int offset,
                         bool wrap,
                         DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
   const int64_t kMaxBlockDim = 512;
   phi::Copy(ctx, x, ctx.GetPlace(), false, out);
 

test/legacy_test/test_tensor_fill_diagonal_.py

@@ -296,5 +296,47 @@ def test_dim_larger2_normal(self):
                 )
 
 
+class TensorFillDiagonal_ZeroSize(unittest.TestCase):
+    def _test_normal(self, shape):
+        expected_np = np.random.random(shape)
+        expected_grad = np.random.random(shape)
+
+        places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not base.core.is_compiled_with_cuda()
+        ):
+            places.append(base.CPUPlace())
+        if base.core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+
+        for idx, p in enumerate(places):
+            if idx == 0:
+                paddle.set_device('cpu')
+            else:
+                paddle.set_device('gpu')
+
+            x = paddle.ones(shape)
+            x.stop_gradient = False
+            y = x * 2
+            y.retain_grads()
+            y.fill_diagonal_(1, offset=0, wrap=True)
+            loss = y.sum()
+            loss.backward()
+
+            self.assertEqual(
+                (y.numpy().astype('float32') == expected_np).all(), True
+            )
+            self.assertEqual(
+                (y.grad.numpy().astype('float32') == expected_grad).all(),
+                True,
+            )
+
+    def test_normal(self):
+        self._test_normal([0, 3])
+        self._test_normal([0, 0])
+
+
 if __name__ == '__main__':
     unittest.main()