【BUPT】【Paddle Tensor 第二期 API支持 0-size Tensor】paddle.matrix_rank 支持 0-size tensor

paddle/phi/kernels/cpu/matrix_rank_tol_kernel.cc

@@ -91,6 +91,27 @@ void MatrixRankTolKernel(const Context& dev_ctx,
   auto dim_out = out->dims();
   int rows = static_cast<int>(dim_x[dim_x.size() - 2]);
   int cols = static_cast<int>(dim_x[dim_x.size() - 1]);
+  PADDLE_ENFORCE_NE(
+      rows,
+      0,
+      errors::InvalidArgument("The input Tensor x's shape[-2] should not "
+                              "be 0, but received %s now.",
+                              dim_x));
+  PADDLE_ENFORCE_NE(
+      cols,
+      0,
+      errors::InvalidArgument("The input Tensor x's shape[-1] should not "
+                              "be 0, but received %s now.",
+                              dim_x));
+  if (x.numel() == 0) {
+    std::vector<int64_t> out_dims_vec(dim_x.size() - 2);
+    for (int i = 0; i < dim_x.size() - 2; ++i) {
+      out_dims_vec[i] = dim_x[i];
+    }
+    out->Resize(phi::make_ddim(out_dims_vec));
+    dev_ctx.template Alloc<int64_t>(out);
+    return;
+  }
   int k = std::min(rows, cols);
   int batches = static_cast<int>(x.numel() / (rows * cols));
 
@@ -179,11 +200,32 @@ void MatrixRankAtolRtolKernel(const Context& dev_ctx,
                               const paddle::optional<DenseTensor>& rtol,
                               bool hermitian,
                               DenseTensor* out) {
-  dev_ctx.template Alloc<int64_t>(out);
   auto dim_x = x.dims();
   auto dim_out = out->dims();
   int rows = static_cast<int>(dim_x[dim_x.size() - 2]);
   int cols = static_cast<int>(dim_x[dim_x.size() - 1]);
+  PADDLE_ENFORCE_NE(
+      rows,
+      0,
+      errors::InvalidArgument("The input Tensor x's shape[-2] should not "
+                              "be 0, but received %s now.",
+                              dim_x));
+  PADDLE_ENFORCE_NE(
+      cols,
+      0,
+      errors::InvalidArgument("The input Tensor x's shape[-1] should not "
+                              "be 0, but received %s now.",
+                              dim_x));
+  if (x.numel() == 0) {
+    std::vector<int64_t> out_dims_vec(dim_x.size() - 2);
+    for (int i = 0; i < dim_x.size() - 2; ++i) {
+      out_dims_vec[i] = dim_x[i];
+    }
+    out->Resize(phi::make_ddim(out_dims_vec));
+    dev_ctx.template Alloc<int64_t>(out);
+    return;
+  }
+  dev_ctx.template Alloc<int64_t>(out);
   int k = std::min(rows, cols);
   int batches = static_cast<int>(x.numel() / (rows * cols));
 

paddle/phi/kernels/gpu/matrix_rank_tol_kernel.cu

@@ -346,14 +346,23 @@ void MatrixRankTolKernel(const Context& dev_ctx,
                     0,
                     common::errors::InvalidArgument(
                         "The input Tensor x's shape[-2] should not "
-                        "be 0, but shape is %s now.",
+                        "be 0, but received %s now.",
                         dim_x));
   PADDLE_ENFORCE_NE(cols,
                     0,
                     common::errors::InvalidArgument(
                         "The input Tensor x's shape[-1] should not "
-                        "be 0, but shape is %s now.",
+                        "be 0, but received %s now.",
                         dim_x));
+  if (x.numel() == 0) {
+    std::vector<int64_t> out_dims_vec(dim_x.size() - 2);
+    for (int i = 0; i < dim_x.size() - 2; ++i) {
+      out_dims_vec[i] = dim_x[i];
+    }
+    out->Resize(phi::make_ddim(out_dims_vec));
+    dev_ctx.template Alloc<int64_t>(out);
+    return;
+  }
   int k = std::min(rows, cols);
   auto numel = x.numel();
   int batches = numel / (rows * cols);
@@ -457,8 +466,6 @@ void MatrixRankAtolRtolKernel(const Context& dev_ctx,
                               bool hermitian,
                               DenseTensor* out) {
   auto* x_data = x.data<T>();
-  dev_ctx.template Alloc<int64_t>(out);
-
   auto dim_x = x.dims();
   auto dim_out = out->dims();
   int rows = dim_x[dim_x.size() - 2];
@@ -467,14 +474,24 @@ void MatrixRankAtolRtolKernel(const Context& dev_ctx,
       rows,
       0,
       errors::InvalidArgument("The input Tensor x's shape[-2] should not "
-                              "be 0, but shape is %s now.",
+                              "be 0, but received %s now.",
                               dim_x));
   PADDLE_ENFORCE_NE(
       cols,
       0,
       errors::InvalidArgument("The input Tensor x's shape[-1] should not "
-                              "be 0, but shape is %s now.",
+                              "be 0, but received %s now.",
                               dim_x));
+  if (x.numel() == 0) {
+    std::vector<int64_t> out_dims_vec(dim_x.size() - 2);
+    for (int i = 0; i < dim_x.size() - 2; ++i) {
+      out_dims_vec[i] = dim_x[i];
+    }
+    out->Resize(phi::make_ddim(out_dims_vec));
+    dev_ctx.template Alloc<int64_t>(out);
+    return;
+  }
+  dev_ctx.template Alloc<int64_t>(out);
   int k = std::min(rows, cols);
   auto numel = x.numel();
   int batches = numel / (rows * cols);

test/legacy_test/test_matrix_rank_atol_rtol_op.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import unittest
 
 import numpy as np
 from op_test import OpTest
+from utils import dygraph_guard, static_guard
 
 import paddle
 from paddle import base, static
@@ -474,5 +476,70 @@ def test_errors(self):
             paddle.linalg.matrix_rank(x, tol=0.2, hermitian=True, rtol=0.2)
 
 
+class TestMatrixRankAtolRtolZeroSizeTensor(unittest.TestCase):
+
+    def _get_places(self):
+        places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not paddle.is_compiled_with_cuda()
+        ):
+            places.append(paddle.CPUPlace())
+        if paddle.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        return places
+
+    def _test_matrix_rank_static(self, place, atol, rtol):
+        with static_guard():
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                x_valid = paddle.static.data(
+                    name='x_valid', shape=[2, 0, 6, 6], dtype='float32'
+                )
+
+                y_valid = paddle.linalg.matrix_rank(
+                    x_valid, atol=atol, rtol=rtol
+                )
+
+                exe = paddle.static.Executor(place)
+                res_valid = exe.run(
+                    feed={'x_valid': np.zeros((2, 0, 6, 6), dtype='float32')},
+                    fetch_list=[y_valid],
+                )
+                self.assertEqual(res_valid[0].shape, tuple(x_valid.shape[:-2]))
+
+    def _test_matrix_rank_dynamic(self, atol, rtol):
+        with dygraph_guard():
+            x_valid = paddle.full((2, 0, 6, 6), 1.0, dtype='float32')
+            x_invalid1 = paddle.full((0, 0), 1.0, dtype='float32')
+            x_invalid2 = paddle.full((2, 3, 0, 0), 1.0, dtype='float32')
+            self.assertRaises(
+                ValueError,
+                paddle.linalg.matrix_rank,
+                x_invalid1,
+                atol=atol,
+                rtol=rtol,
+            )
+            self.assertRaises(
+                ValueError,
+                paddle.linalg.matrix_rank,
+                x_invalid2,
+                atol=atol,
+                rtol=rtol,
+            )
+
+            y_valid = paddle.linalg.matrix_rank(x_valid, atol=atol, rtol=rtol)
+            self.assertEqual(y_valid.shape, x_valid.shape[:-2])
+
+    def test_matrix_rank_tensor(self):
+        atol = 0.2
+        rtol = 0.2
+        for place in [paddle.CPUPlace()]:
+            self._test_matrix_rank_static(place, atol, rtol)
+        self._test_matrix_rank_dynamic(atol, rtol)
+
+
 if __name__ == '__main__':
     unittest.main()

test/legacy_test/test_matrix_rank_op.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 from op_test import OpTest
+from utils import dygraph_guard, static_guard
 
 import paddle
 from paddle import base, static
@@ -230,5 +231,67 @@ def test_static(self):
                 np.testing.assert_allclose(fetches[0], rank_np, rtol=1e-05)
 
 
+class TestMatrixRankZeroSizeTensor(unittest.TestCase):
+
+    def _get_places(self):
+        places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not paddle.is_compiled_with_cuda()
+        ):
+            places.append(paddle.CPUPlace())
+        if paddle.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        return places
+
+    def _test_matrix_rank_static(self, place):
+        with static_guard():
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                x_valid = paddle.static.data(
+                    name='x_valid', shape=[2, 0, 6, 6], dtype='float32'
+                )
+
+                y_valid = paddle.linalg.matrix_rank(x_valid)
+
+                exe = paddle.static.Executor(place)
+                res_valid = exe.run(
+                    feed={'x_valid': np.zeros((2, 0, 6, 6), dtype='float32')},
+                    fetch_list=[y_valid],
+                )
+                self.assertEqual(res_valid[0].shape, tuple(x_valid.shape[:-2]))
+
+    def _test_matrix_rank_dynamic_cpu(self):
+        with dygraph_guard():
+            paddle.set_device("cpu")
+            x_valid = paddle.full((2, 0, 6, 6), 1.0, dtype='float32')
+            x_invalid1 = paddle.full((0, 0), 1.0, dtype='float32')
+            x_invalid2 = paddle.full((2, 3, 0, 0), 1.0, dtype='float32')
+            self.assertRaises(ValueError, paddle.linalg.matrix_rank, x_invalid1)
+            self.assertRaises(ValueError, paddle.linalg.matrix_rank, x_invalid2)
+
+            y_valid = paddle.linalg.matrix_rank(x_valid)
+            self.assertEqual(y_valid.shape, x_valid.shape[:-2])
+
+    def _test_matrix_rank_dynamic_gpu(self):
+        with dygraph_guard():
+            x_valid = paddle.full((2, 0, 6, 6), 1.0, dtype='float32')
+            x_invalid1 = paddle.full((0, 0), 1.0, dtype='float32')
+            x_invalid2 = paddle.full((2, 3, 0, 0), 1.0, dtype='float32')
+            self.assertRaises(ValueError, paddle.linalg.matrix_rank, x_invalid1)
+            self.assertRaises(ValueError, paddle.linalg.matrix_rank, x_invalid2)
+
+            y_valid = paddle.linalg.matrix_rank(x_valid)
+            self.assertEqual(y_valid.shape, x_valid.shape[:-2])
+
+    def test_matrix_rank_tensor(self):
+        for place in self._get_places():
+            self._test_matrix_rank_static(place)
+        self._test_matrix_rank_dynamic_cpu()
+        self._test_matrix_rank_dynamic_gpu()
+
+
 if __name__ == '__main__':
     unittest.main()