[0-size suppor] fix the eighGrad and eigvalshGrad (#72756)

DanielSun11 · web-flow · commit dc5e99fda58a · 2025-05-20T15:38:35.000+08:00
diff --git a/paddle/phi/kernels/impl/eigh_grad_kernel_impl.h b/paddle/phi/kernels/impl/eigh_grad_kernel_impl.h
@@ -36,6 +36,9 @@ void EighGradKernel(const Context& dev_ctx,
                     const DenseTensor& dout_v,
                     DenseTensor* dx) {
   dev_ctx.template Alloc<T>(dx);
+  if (out_v.numel() == 0) {
+    return;
+  }
   auto& dims = out_v.dims();
   const int m = dims[dims.size() - 1];
   DenseTensor tV =
diff --git a/paddle/phi/kernels/impl/eigvalsh_grad_kernel_impl.h b/paddle/phi/kernels/impl/eigvalsh_grad_kernel_impl.h
@@ -30,6 +30,10 @@ void EigvalshGradKernel(const Context& dev_ctx,
                         const std::string& uplo UNUSED,
                         bool is_test UNUSED,
                         DenseTensor* x_grad) {
+  if (x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   auto tV = phi::TransposeLast2Dim<T>(dev_ctx, phi::Conj<T>(dev_ctx, out_v));
 
   x_grad->Resize(out_v.dims());
diff --git a/test/legacy_test/test_eigh_op.py b/test/legacy_test/test_eigh_op.py
@@ -295,6 +295,8 @@ def setUp(self):
             if paddle.is_compiled_with_cuda()
             else paddle.CPUPlace()
         )
+        self.rtol = 1e-5  # for test_eigh_grad
+        self.atol = 1e-5  # for test_eigh_grad
         np.random.seed(123)
 
     def init_input_shape(self):
@@ -345,11 +347,33 @@ def test_in_dynamic_mode(self):
                 self.real_data, actual_w.numpy(), actual_v.numpy()
             )
 
+    def test_eigh_grad(self):
+        paddle.disable_static()
+        x = paddle.to_tensor(self.real_data, stop_gradient=False)
+        w, v = paddle.linalg.eigh(x)
+        self.trans_dims = [
+            *range(len(self.x_shape) - 2),
+            len(self.x_shape) - 1,
+            len(self.x_shape) - 2,
+        ]
+        (w.sum() + paddle.abs(v).sum()).backward()
+        np.testing.assert_allclose(
+            abs(x.grad.numpy()),
+            abs(x.grad.numpy().conj().transpose(self.trans_dims)),
+            rtol=self.rtol,
+            atol=self.atol,
+        )
+
 
 class TestEighBatchAPIZeroSize(TestEighAPIZeroSize):
     def init_input_shape(self):
         self.x_shape = [0, 5, 5]
 
 
+class TestEighBatchAPIZeroSize1(TestEighAPIZeroSize):
+    def init_input_shape(self):
+        self.x_shape = [5, 0, 0]
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/legacy_test/test_eigvalsh_op.py b/test/legacy_test/test_eigvalsh_op.py
@@ -268,6 +268,8 @@ def setUp(self):
         np.random.seed(123)
         self.init_input_shape()
         self.init_input_data()
+        self.rtol = 1e-5  # for test_eigh_grad
+        self.atol = 1e-5  # for test_eigh_grad
 
     def init_input_shape(self):
         self.x_shape = [0, 0]
@@ -318,11 +320,39 @@ def test_in_dynamic_mode(self):
             actual_w = paddle.linalg.eigvalsh(input_real_data)
             compare_shape_result(actual_w.numpy(), expected_w)
 
+    def test_eigvalsh_grad(self):
+        paddle.disable_static(self.place)
+        self.trans_dims = [
+            *range(len(self.x_shape) - 2),
+            len(self.x_shape) - 1,
+            len(self.x_shape) - 2,
+        ]
+        x = paddle.to_tensor(self.real_data, stop_gradient=False)
+        w = paddle.linalg.eigvalsh(x)
+        (w.sum()).backward()
+
+        # compare with eigh
+        y = paddle.to_tensor(self.real_data, stop_gradient=False)
+        y_v, y_w = paddle.linalg.eigh(y)
+        (y_v.sum()).backward()
 
-class TestEigvalshBatchAPIZeroSize(TestEigvalshAPIZeroSize):
+        np.testing.assert_allclose(
+            abs(x.grad.numpy()),
+            abs(y.grad.numpy()),
+            rtol=self.rtol,
+            atol=self.atol,
+        )
+
+
+class TestEigvalshBatchAPIZeroSize1(TestEigvalshAPIZeroSize):
     def init_input_shape(self):
         self.x_shape = [0, 5, 5]
 
 
+class TestEigvalshBatchAPIZeroSize2(TestEigvalshAPIZeroSize):
+    def init_input_shape(self):
+        self.x_shape = [5, 0, 0]
+
+
 if __name__ == "__main__":
     unittest.main()
