【BUPT】【Paddle Tensor 第二期 API支持 0-size Tensor】paddle.matrix_power 支持 0-size tensor (#70098)

ZHOU05030 · web-flow · commit 8e03e869c2ea · 2024-12-12T19:35:14.000+08:00
* add matrix_power 0 size

* fix

* fix bug

* fix bug

* fix

* 添加测试情况

* fix

* fix
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -2427,10 +2427,11 @@ void MatrixPowerInferMeta(const MetaTensor& x, int n, MetaTensor* out) {
                         "received a %d dimension tensor.",
                         n_dim));
   for (int i = 0; i < n_dim; ++i)
-    PADDLE_ENFORCE_NE(dims[i],
-                      0,
-                      common::errors::InvalidArgument(
-                          "The size of Input(X) should not be 0."));
+    PADDLE_ENFORCE_GE(
+        dims[i],
+        0,
+        common::errors::InvalidArgument(
+            "The size of Input(X) should greater than or equal to 0."));
   PADDLE_ENFORCE_EQ(dims[n_dim - 2],
                     dims[n_dim - 1],
                     common::errors::InvalidArgument(
diff --git a/paddle/phi/kernels/impl/matrix_power_kernel_impl.h b/paddle/phi/kernels/impl/matrix_power_kernel_impl.h
@@ -196,6 +196,11 @@ void MatrixPowerKernel(const Context& ctx,
           "X's shape[-2] = %d and shape[-1] = %d.",
           x_dims[x_ndim - 2],
           x_dims[x_ndim - 1]));
+  if (x.numel() == 0) {
+    Out->Resize(X->dims());
+    ctx.template Alloc<T>(Out);
+    return;
+  }
 
   MatrixPowerFunction<Context, T>(X, n, Out, ctx);
 }
diff --git a/test/legacy_test/test_matrix_power_op.py b/test/legacy_test/test_matrix_power_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
@@ -333,20 +334,6 @@ def test_errors(self):
         )
         self.assertRaises(ValueError, paddle.linalg.matrix_power, input, 2)
 
-        # The size of input should not be 0
-        input = paddle.static.data(
-            name="input_4", shape=[1, 1, 0, 0], dtype="float32"
-        )
-        self.assertRaises(ValueError, paddle.linalg.matrix_power, input, 2)
-
-        # The size of input should not be 0
-        input = paddle.static.data(
-            name="input_5", shape=[0, 0], dtype="float32"
-        )
-        self.assertRaises(
-            ValueError, paddle.linalg.matrix_power, input, -956301312
-        )
-
     def test_old_ir_errors(self):
         if paddle.framework.use_pir_api():
             return
@@ -409,6 +396,78 @@ def test_dygraph(self):
                     print("The mat is singular")
 
 
+class TestMatrixPowerEmptyTensor(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(base.CPUPlace())
+        if paddle.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        return places
+
+    def _test_matrix_power_empty_static(self, place):
+        with static_guard():
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                x2 = paddle.static.data(
+                    name='x2', shape=[0, 6], dtype='float32'
+                )
+                x3 = paddle.static.data(
+                    name='x3', shape=[6, 0], dtype='float32'
+                )
+                x4 = paddle.static.data(
+                    name='x4', shape=[0, 0, 2, 3], dtype='float32'
+                )
+                self.assertRaises(TypeError, paddle.linalg.matrix_power, x2)
+                self.assertRaises(TypeError, paddle.linalg.matrix_power, x3)
+                self.assertRaises(TypeError, paddle.linalg.matrix_power, x4)
+
+                x = paddle.static.data(name='x', shape=[0, 0], dtype='float32')
+                y = paddle.linalg.matrix_power(x, 2)
+                x5 = paddle.static.data(
+                    name='x5', shape=[2, 3, 0, 0], dtype='float32'
+                )
+                y5 = paddle.linalg.matrix_power(x5, 2)
+                exe = paddle.static.Executor(place)
+                res = exe.run(
+                    feed={
+                        'x2': np.zeros((0, 6), dtype='float32'),
+                        'x3': np.zeros((6, 0), dtype='float32'),
+                        'x4': np.zeros((0, 0, 2, 3), dtype='float32'),
+                        'x': np.zeros((0, 0), dtype='float32'),
+                        'x5': np.zeros((2, 3, 0, 0), dtype='float32'),
+                    },
+                    fetch_list=[y, y5],
+                )
+                self.assertEqual(res[0].shape, (0, 0))
+                self.assertEqual(res[1].shape, (2, 3, 0, 0))
+
+    def _test_matrix_power_empty_dynamtic(self):
+        with dygraph_guard():
+            x2 = paddle.full((0, 6), 1.0, dtype='float32')
+            x3 = paddle.full((6, 0), 1.0, dtype='float32')
+            x4 = paddle.full((2, 3, 0, 0), 1.0, dtype='float32')
+            x5 = paddle.full((0, 0, 2, 3), 1.0, dtype='float32')
+            self.assertRaises(TypeError, paddle.linalg.matrix_power, x2)
+            self.assertRaises(TypeError, paddle.linalg.matrix_power, x3)
+            self.assertRaises(TypeError, paddle.linalg.matrix_power, x5)
+            x = paddle.full((0, 0), 1.0, dtype='float32')
+            y = paddle.linalg.matrix_power(x, 2)
+            y4 = paddle.linalg.matrix_power(x4, 2)
+            self.assertEqual(y4.shape, [2, 3, 0, 0])
+            self.assertEqual(y.shape, [0, 0])
+
+    def test_matrix_power_empty_tensor(self):
+        for place in self._get_places():
+            self._test_matrix_power_empty_static(place)
+        self._test_matrix_power_empty_dynamtic()
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()
