[0-size Tensor No. 103、292] Add 0-size Tensor support for paddle.linalg.cholesky_solve API. (#72943)

DanielSun11 · web-flow · commit 614aab2884f7 · 2025-05-28T11:08:14.000+08:00
* fix 0 size Tensor for expand kernel in onednn

* format

* expand_grad support 0 size Tensor

* fix bug and add  unittest

* modify the unittest for onednn kernel

* modify expand unittest to force enable onednn op

* fix cholesky_sovle
diff --git a/paddle/phi/kernels/impl/cholesky_solve_grad_kernel_impl.h b/paddle/phi/kernels/impl/cholesky_solve_grad_kernel_impl.h
@@ -20,14 +20,14 @@
 #include "paddle/phi/kernels/elementwise_add_kernel.h"
 #include "paddle/phi/kernels/empty_kernel.h"
 #include "paddle/phi/kernels/expand_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/blas/blas.h"
 #include "paddle/phi/kernels/funcs/common_shape.h"
 #include "paddle/phi/kernels/funcs/complex_functors.h"
 #include "paddle/phi/kernels/funcs/for_range.h"
 #include "paddle/phi/kernels/funcs/matrix_reduce.h"
 #include "paddle/phi/kernels/funcs/tril_triu_compute.h"
 #include "paddle/phi/kernels/transpose_kernel.h"
-
 namespace phi {
 
 template <typename T, typename Context>
@@ -39,6 +39,23 @@ void CholeskySolveGradKernel(const Context& dev_ctx,
                              bool upper,
                              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;
+  }
   // get broadcast dim
   std::vector<int64_t> x_bst_dims_vec;
   std::vector<int64_t> y_bst_dims_vec;
diff --git a/paddle/phi/kernels/impl/cholesky_solve_kernel_impl.h b/paddle/phi/kernels/impl/cholesky_solve_kernel_impl.h
@@ -42,6 +42,10 @@ void CholeskySolveKernel(const Context& dev_ctx,
                          const DenseTensor& y,
                          bool upper,
                          DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   // get broadcast dim
   std::vector<int64_t> x_bst_dims_vec;
   std::vector<int64_t> y_bst_dims_vec;
diff --git a/test/legacy_test/test_cholesky_solve_op.py b/test/legacy_test/test_cholesky_solve_op.py
@@ -282,5 +282,63 @@ def test_errors_2(self):
             self.assertRaises(ValueError, paddle.linalg.cholesky_solve, x7, y7)
 
 
+# API function test
+class TestCholeskySolveAPIZeroSize(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2025)
+        self.place = [paddle.CPUPlace()]
+        self.dtype = "float64"
+        self.upper = True
+        if core.is_compiled_with_cuda():
+            self.place.append(paddle.CUDAPlace(0))
+        self.init_shape()
+
+    def init_shape(self):
+        self.x_shape = [10, 0]
+        self.y_shape = [10, 10]
+        self.expected_shape = [10, 0]
+        # test in dynamic mode
+
+    def test_dygraph(self):
+        def run(place):
+            paddle.disable_static(place)
+            x_np = np.random.random(self.x_shape).astype(self.dtype)
+            y_np = np.random.random(self.y_shape).astype(self.dtype)
+
+            x = paddle.to_tensor(x_np, stop_gradient=False)
+            y = paddle.to_tensor(y_np, stop_gradient=False)
+            z = paddle.linalg.cholesky_solve(x, y, upper=self.upper)
+            loss = paddle.sum(z)
+            loss.backward()
+
+            self.assertEqual(z.shape, self.expected_shape)
+            self.assertEqual(x.shape, x.grad.shape)
+            self.assertEqual(y.shape, y.grad.shape)
+
+        for idx, place in enumerate(self.place):
+            run(place)
+
+
+class TestCholeskySolveAPIZeroSize1(TestCholeskySolveAPIZeroSize):
+    def init_shape(self):
+        self.x_shape = [0, 6]
+        self.y_shape = [0, 0]
+        self.expected_shape = [0, 6]
+
+
+class TestCholeskySolveAPIZeroSize2(TestCholeskySolveAPIZeroSize):
+    def init_shape(self):
+        self.x_shape = [1, 10, 6]
+        self.y_shape = [0, 10, 10]
+        self.expected_shape = [0, 10, 6]
+
+
+class TestCholeskySolveAPIZeroSize3(TestCholeskySolveAPIZeroSize):
+    def init_shape(self):
+        self.x_shape = [0, 0, 0]
+        self.y_shape = [0, 0, 0]
+        self.expected_shape = [0, 0, 0]
+
+
 if __name__ == "__main__":
     unittest.main()
