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

* fix

* update

* fix

* fix

* fix

Author: aquagull

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

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import unittest
 
 import numpy as np
@@ -64,14 +63,7 @@ def test_check_output(self):
         self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        places = []
-        if (
-            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
-            in ['1', 'true', 'on']
-            or not core.is_compiled_with_cuda()
-            or core.is_compiled_with_rocm()
-        ):
-            places.append(base.CPUPlace())
+        places = [base.CPUPlace()]
         if core.is_compiled_with_cuda() and (not core.is_compiled_with_rocm()):
             places.append(base.CUDAPlace(0))
         for p in places:
@@ -161,6 +153,11 @@ def init_config(self):
         self._input_shape = (32, 32)
 
 
+class TestCholeskyOpZeroSize(TestCholeskyOp):
+    def init_config(self):
+        self._input_shape = (0, 0)
+
+
 class TestDygraph(unittest.TestCase):
     def test_dygraph(self):
         if core.is_compiled_with_rocm():
@@ -176,27 +173,20 @@ def test_dygraph(self):
 
 class TestCholeskySingularAPI(unittest.TestCase):
     def setUp(self):
-        self.places = []
-        if (
-            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
-            in ['1', 'true', 'on']
-            or not core.is_compiled_with_cuda()
-            or core.is_compiled_with_rocm()
-        ):
-            self.places.append(base.CPUPlace())
+        self.places = [base.CPUPlace()]
         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 +201,9 @@ 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])
+            self.check_static_result(place=place, input_shape=[5, 0, 0])
 
     def test_dygraph(self):
         for place in self.places:
@@ -222,9 +214,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: