fix solve grad (#72806)

DanielSun11 · web-flow · commit 26992ec7f3e3 · 2025-05-21T11:18:55.000+08:00
diff --git a/paddle/phi/kernels/impl/solve_grad_kernel_impl.h b/paddle/phi/kernels/impl/solve_grad_kernel_impl.h
@@ -17,6 +17,7 @@ limitations under the License. */
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/kernels/expand_as_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/blas/blas.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/funcs/matrix_solve.h"
@@ -78,6 +79,24 @@ void SolveGradKernel(const Context& dev_ctx,
                      const DenseTensor& dout,
                      DenseTensor* dx,
                      DenseTensor* dy) {
+  if (dout.numel() == 0) {
+    if (dx) {
+      dev_ctx.template Alloc<T>(dx);
+      if (dx->numel() != 0) {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dx->dims())), 0, dx);
+      }
+    }
+    if (dy) {
+      dev_ctx.template Alloc<T>(dy);
+      if (dy->numel() != 0) {
+        phi::Full<T, Context>(
+            dev_ctx, phi::IntArray(common::vectorize(dy->dims())), 0, dy);
+      }
+    }
+    return;
+  }
+
   bool is_vector = false;
   is_vector = is_vector_rhs(x, y);
   DenseTensor tmp_y;
diff --git a/test/legacy_test/test_solve_op.py b/test/legacy_test/test_solve_op.py
@@ -985,22 +985,39 @@ def run(place, x_shape, y_shape):
                 input_x_np = np.random.random(x_shape).astype(self.dtype)
                 input_y_np = np.random.random(y_shape).astype(self.dtype)
 
-                tensor_input_x = paddle.to_tensor(input_x_np)
-                tensor_input_y = paddle.to_tensor(input_y_np)
+                tensor_input_x = paddle.to_tensor(
+                    input_x_np, stop_gradient=False
+                )
+                tensor_input_y = paddle.to_tensor(
+                    input_y_np, stop_gradient=False
+                )
 
                 numpy_output = np.linalg.solve(input_x_np, input_y_np)
                 paddle_output = paddle.linalg.solve(
-                    tensor_input_x, tensor_input_y, left=False
+                    tensor_input_x, tensor_input_y, left=True
                 )
                 np.testing.assert_allclose(
                     numpy_output, paddle_output.numpy(), rtol=0.0001
                 )
                 self.assertEqual(
                     numpy_output.shape, paddle_output.numpy().shape
                 )
+                loss = paddle.sum(paddle_output)
+                loss.backward()
+                np.testing.assert_allclose(
+                    tensor_input_x.grad.shape, tensor_input_x.shape
+                )
+                np.testing.assert_allclose(
+                    tensor_input_y.grad.shape, tensor_input_y.shape
+                )
 
         for place in self.place:
+            run(place, x_shape=[1, 10, 10], y_shape=[1, 10, 10])
+            run(place, x_shape=[0, 10, 10], y_shape=[0, 10, 10])
+            run(place, x_shape=[0, 10, 10], y_shape=[1, 10, 10])
             run(place, x_shape=[10, 0, 0], y_shape=[10, 0, 0])
+            run(place, x_shape=[10, 1, 1], y_shape=[10, 1, 0])
+
             with self.assertRaises(ValueError) as context:
                 run(place, x_shape=[10, 0, 0], y_shape=[10])
 
