[API] Support tensor shape in reshape with compatible API (PaddlePaddle#76025)

Author: SigureMo

python/paddle/utils/decorator_utils.py

@@ -21,13 +21,23 @@
 
 from typing_extensions import ParamSpec
 
+import paddle
+
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
 _InputT = ParamSpec("_InputT")
 _RetT = TypeVar("_RetT")
 
 
+def _is_in_or_scalar_tensor(x):
+    if isinstance(x, int):
+        return True
+    if isinstance(x, (paddle.Tensor, paddle.pir.Value)):
+        return x.ndim == 0
+    return False
+
+
 class DecoratorBase:
     """Decorative base class, providing a universal decorative framework.
 
@@ -410,8 +420,8 @@ def wrapper(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
                 kwargs["shape_or_dtype"] = kwargs.pop("dtype")
             elif ("size" in kwargs) and ("shape_or_dtype" not in kwargs):
                 kwargs["shape_or_dtype"] = kwargs.pop("size")
-            elif len(args) >= 2 and type(args[1]) is int:
-                if all(type(arg) is int for arg in args[1:]):
+            elif len(args) >= 2 and _is_in_or_scalar_tensor(args[1]):
+                if all(_is_in_or_scalar_tensor(arg) for arg in args[1:]):
                     kwargs["x"] = args[0]
                     kwargs['shape_or_dtype'] = list(args[1:])
                     args = ()
@@ -542,8 +552,8 @@ def decorator(func: Callable[_InputT, _RetT]) -> Callable[_InputT, _RetT]:
         def wrapper(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
             if ("input" in kwargs) and ("x" not in kwargs):
                 kwargs["x"] = kwargs.pop("input")
-            elif len(args) >= 2 and type(args[1]) is int:
-                if all(type(arg) is int for arg in args[1:]):
+            elif len(args) >= 2 and _is_in_or_scalar_tensor(args[1]):
+                if all(_is_in_or_scalar_tensor(arg) for arg in args[1:]):
                     kwargs["x"] = args[0]
                     kwargs['shape'] = list(args[1:])
                     args = ()
@@ -614,8 +624,8 @@ def wrapper(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
                 kwargs["x"] = kwargs.pop("input")
             if ("size" in kwargs) and ("shape" not in kwargs):
                 kwargs["shape"] = kwargs.pop("size")
-            elif len(args) >= 2 and type(args[1]) is int:
-                if all(type(arg) is int for arg in args[1:]):
+            elif len(args) >= 2 and _is_in_or_scalar_tensor(args[1]):
+                if all(_is_in_or_scalar_tensor(arg) for arg in args[1:]):
                     kwargs["x"] = args[0]
                     kwargs['shape'] = list(args[1:])
                     args = ()

test/legacy_test/test_reshape_op.py

@@ -23,6 +23,7 @@
     is_custom_device,
     skip_check_grad_ci,
 )
+from utils import dygraph_guard, static_guard
 
 import paddle
 from paddle import base
@@ -942,6 +943,134 @@ def run_test_cases(place):
                 run_test_cases(get_device_place())
 
 
+class TestReshapeWithTensorShape(unittest.TestCase):
+    """
+    reshape supports shape like:
+    paddle.reshape(x, shape=[1, 2, 3])
+    paddle.reshape(x, shape=[1, Tensor(2), 3])
+    paddle.reshape(x, shape=Tensor([1, 2, 3]))
+    paddle.reshape(x, 1, 2, 3)  # Compatible usage
+    paddle.reshape(x, 1, Tensor(2), 3)  # Compatible usage
+    """
+
+    @static_guard()
+    def check_reshape_static(
+        self, fn, x_shape, expected_out_shape, dynamic_dims=[]
+    ):
+        main_program = Program()
+        with program_guard(main_program):
+            x = paddle.static.data('x', shape=x_shape, dtype='float32')
+            out = fn(x)
+            if dynamic_dims:
+                expected_out_shape_with_dynamic = list(expected_out_shape)
+                for dim in dynamic_dims:
+                    expected_out_shape_with_dynamic[dim] = -1
+                self.assertEqual(out.shape, expected_out_shape_with_dynamic)
+            else:
+                self.assertEqual(out.shape, expected_out_shape)
+
+        exe = paddle.static.Executor()
+        (out_np,) = exe.run(
+            main_program,
+            feed={'x': np.random.random(x_shape)},
+            fetch_list=[out],
+        )
+        self.assertEqual(list(out_np.shape), expected_out_shape)
+
+    @dygraph_guard()
+    def check_reshape_dygraph(self, fn, x_shape, expected_out_shape):
+        x = paddle.to_tensor(np.random.random(x_shape).astype('float32'))
+        out = fn(x)
+        self.assertEqual(list(out.shape), expected_out_shape)
+
+    def check_reshape(self, fn, x_shape, expected_out_shape):
+        self.check_reshape_static(fn, x_shape, expected_out_shape)
+        self.check_reshape_dygraph(fn, x_shape, expected_out_shape)
+
+    def test_reshape_with_list_int(self):
+        def reshape_fn(x):
+            return paddle.reshape(x, shape=[2, 3, 4])
+
+        self.check_reshape(reshape_fn, [2, 12], [2, 3, 4])
+
+    def test_reshape_with_list_scalar_tensor(self):
+        def reshape_fn(x):
+            dim0 = paddle.full([], 2, dtype='int64')
+            dim1 = paddle.full([], 3, dtype='int64')
+            dim2 = paddle.full([], 4, dtype='int64')
+            return paddle.reshape(x, shape=[dim0, dim1, dim2])
+
+        self.check_reshape(reshape_fn, [2, 12], [2, 3, 4])
+
+    def test_reshape_with_list_scalar_tensor_dynamic_dim(self):
+        def reshape_fn(x):
+            dim0 = paddle.full([], 1, dtype='int64') + 1  # dynamic dim
+            dim1 = paddle.full([], 3, dtype='int64')
+            dim2 = paddle.full([], 4, dtype='int64')
+            return paddle.reshape(x, shape=[dim0, dim1, dim2])
+
+        self.check_reshape_static(
+            reshape_fn,
+            x_shape=[2, 12],
+            expected_out_shape=[2, 3, 4],
+            dynamic_dims=[0],
+        )
+
+    def test_reshape_with_list_mix_int_tensor(self):
+        def reshape_fn(x):
+            dim1 = paddle.full([], 3, dtype='int64')
+            return paddle.reshape(x, shape=[2, dim1, 4])
+
+        self.check_reshape(reshape_fn, [2, 12], [2, 3, 4])
+
+    def test_reshape_with_tensor_dynamic_dim(self):
+        def reshape_fn(x):
+            shape_tensor = paddle.to_tensor([1, 2, 3]) + 1  # all dynamic dims
+            return paddle.reshape(x, shape=shape_tensor)
+
+        self.check_reshape_static(
+            reshape_fn,
+            x_shape=[2, 12],
+            expected_out_shape=[2, 3, 4],
+            dynamic_dims=[0, 1, 2],
+        )
+
+    def test_reshape_with_tensor(self):
+        def reshape_fn(x):
+            shape_tensor = paddle.stack(
+                [
+                    paddle.full([], 2, dtype='int64'),
+                    paddle.full([], 3, dtype='int64'),
+                    paddle.full([], 4, dtype='int64'),
+                ]
+            )
+            return paddle.reshape(x, shape=shape_tensor)
+
+        self.check_reshape(reshape_fn, [2, 12], [2, 3, 4])
+
+    def test_reshape_with_list_int_compatible(self):
+        def reshape_fn(x):
+            return paddle.reshape(x, 2, 3, 4)
+
+        self.check_reshape(reshape_fn, [2, 12], [2, 3, 4])
+
+    def test_reshape_with_list_scalar_tensor_compatible(self):
+        def reshape_fn(x):
+            dim0 = paddle.full([], 2, dtype='int64')
+            dim1 = paddle.full([], 3, dtype='int64')
+            dim2 = paddle.full([], 4, dtype='int64')
+            return paddle.reshape(x, dim0, dim1, dim2)
+
+        self.check_reshape(reshape_fn, [2, 12], [2, 3, 4])
+
+    def test_reshape_with_list_mix_int_tensor_compatible(self):
+        def reshape_fn(x):
+            dim1 = paddle.full([], 3, dtype='int64')
+            return paddle.reshape(x, 2, dim1, 4)
+
+        self.check_reshape(reshape_fn, [2, 12], [2, 3, 4])
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()