【 Paddle Tensor 规范化第二期 】paddle.linalg.cholesky适配0-size Tensor

paddle/phi/kernels/cpu/cholesky_kernel.cc

@@ -33,6 +33,11 @@ void CholeskyKernel(const Context& dev_ctx,
   using OutputMatrixMap = Eigen::Map<EigenMatrix>;
 
   auto& dims = x.dims();
+  if (x.numel() == 0) {
+    out->Resize(dims);
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   int batch_count = 1;
   for (int i = 0; i < dims.size() - 2; i++) {
     batch_count *= static_cast<int>(dims[i]);

paddle/phi/kernels/gpu/cholesky_kernel.cu

@@ -115,6 +115,11 @@ void CholeskyKernel(const Context& dev_ctx,
                     bool upper,
                     DenseTensor* out) {
   auto& dims = x.dims();
+  if (x.numel() == 0) {
+    out->Resize(dims);
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   int batch_count = 1;
   for (int i = 0; i < dims.size() - 2; i++) {
     batch_count *= dims[i];

paddle/phi/kernels/impl/cholesky_grad_kernel_impl.h

@@ -245,6 +245,11 @@ void CholeskyGradKernel(const Context& dev_ctx,
   auto* x_grad_data = dev_ctx.template Alloc<T>(x_grad);
 
   auto& dims = out.dims();
+  if (out.numel() == 0) {
+    x_grad->Resize(dims);
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   int batch_count = 1;
   for (int i = 0; i < dims.size() - 2; i++) {
     batch_count *= dims[i];

test/legacy_test/test_cholesky_op.py

@@ -161,6 +161,11 @@ def init_config(self):
         self._input_shape = (32, 32)
 
 
+class TestCholeskyOpZeroSize(TestCholeskyOp):
+    def init_config(self):
+        self._input_shape = (0, 2)
+
+
 class TestDygraph(unittest.TestCase):
     def test_dygraph(self):
         if core.is_compiled_with_rocm():
@@ -187,16 +192,16 @@ def setUp(self):
         if core.is_compiled_with_cuda() and (not core.is_compiled_with_rocm()):
             self.places.append(base.CUDAPlace(0))
 
-    def check_static_result(self, place, with_out=False):
+    def check_static_result(self, place, input_shape, with_out=False):
         with paddle.static.program_guard(
             paddle.static.Program(), paddle.static.Program()
         ):
             input = paddle.static.data(
-                name="input", shape=[4, 4], dtype="float64"
+                name="input", shape=input_shape, dtype="float64"
             )
             result = paddle.cholesky(input)
 
-            input_np = np.zeros([4, 4]).astype("float64")
+            input_np = np.zeros(input_shape).astype("float64")
 
             exe = base.Executor(place)
             try:
@@ -211,7 +216,8 @@ def check_static_result(self, place, with_out=False):
 
     def test_static(self):
         for place in self.places:
-            self.check_static_result(place=place)
+            self.check_static_result(place=place, input_shape=[4, 4])
+            self.check_static_result(place=place, input_shape=[0, 0])
 
     def test_dygraph(self):
         for place in self.places:
@@ -222,9 +228,12 @@ def test_dygraph(self):
                         [[10, 11, 12], [13, 14, 15], [16, 17, 18]],
                     ]
                 ).astype("float64")
+                input_np_zero = np.zeros((0, 3, 3), dtype="float64")
                 input = paddle.to_tensor(input_np)
+                input_zero = paddle.to_tensor(input_np_zero)
                 try:
                     result = paddle.cholesky(input)
+                    result_zero = paddle.cholesky(input_zero)
                 except RuntimeError as ex:
                     print("The mat is singular")
                 except ValueError as ex: