some create api support more usage (PaddlePaddle#74494)

Author: zhwesky2010

python/paddle/tensor/creation.py

@@ -24,7 +24,7 @@
 
 import paddle
 from paddle import _C_ops
-from paddle.utils.decorator_utils import ParamAliasDecorator
+from paddle.utils.decorator_utils import ParamAliasDecorator, SizeArgsDecorator
 from paddle.utils.inplace_utils import inplace_apis_in_dygraph_only
 
 from ..base.data_feeder import (
@@ -1241,6 +1241,7 @@ def fill_constant(
         return out
 
 
+@SizeArgsDecorator()
 def ones(
     shape: ShapeLike, dtype: DTypeLike | None = None, name: str | None = None
 ) -> paddle.Tensor:

python/paddle/utils/decorator_utils.py

@@ -91,30 +91,27 @@ def process(
         return args, processed_kwargs
 
 
-class ForbidKeywordsDecorator(DecoratorBase):
-    """A decorator that hints users to use the correct `compat` functions, when erroneous keyword arguments are detected"""
+# *size => shape decorator
+class SizeArgsDecorator(DecoratorBase):
+    """
+    Usage Example:
 
-    def __init__(
-        self, illegal_keys: list[str], func_name: str, correct_name: str
-    ) -> None:
-        super().__init__()
-        self.illegal_keys = (
-            [illegal_keys] if isinstance(illegal_keys, str) else illegal_keys
-        )
-        self.func_name = func_name
-        self.correct_name = correct_name
+    paddle.ones(1, dtype=paddle.float32)
+    paddle.ones(1, 2, 3, dtype=paddle.float32)
+    paddle.ones([1, 2, 3], dtype=paddle.float32)
+    paddle.ones(size=[1, 2, 3], dtype=paddle.float32)
+
+    paddle.ones([1, 2, 3], paddle.float32)
+    paddle.ones(shape=[1, 2, 3], dtype=paddle.float32)
+    """
 
     def process(
         self, args: tuple[Any, ...], kwargs: dict[str, Any]
     ) -> tuple[tuple[Any, ...], dict[str, Any]]:
-        found_keys = [key for key in self.illegal_keys if key in kwargs]
-
-        if found_keys:
-            keys_str = ", ".join(f"'{key}'" for key in found_keys)
-            plural = "s" if len(found_keys) > 1 else ""
+        if 'size' in kwargs:
+            kwargs['shape'] = kwargs.pop('size')
+        elif len(args) >= 1 and isinstance(args[0], int):
+            kwargs['shape'] = list(args)
+            args = ()
 
-            raise TypeError(
-                f"{self.func_name}() received unexpected keyword argument{plural} {keys_str}. "
-                f"\nDid you mean to use {self.correct_name}() instead?"
-            )
         return args, kwargs

test/legacy_test/test_ones_op.py

@@ -20,38 +20,121 @@
 
 
 class ApiOnesTest(unittest.TestCase):
-    def test_paddle_ones(self):
+    def test_static_ones(self):
+        paddle.enable_static()
+        with paddle.static.program_guard(paddle.static.Program()):
+            ones = paddle.ones(10, dtype=paddle.float32)
+            place = paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            (result,) = exe.run(fetch_list=[ones])
+            expect = np.ones([10], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
+        with paddle.static.program_guard(paddle.static.Program()):
+            ones = paddle.ones(10, 2, 3, dtype=paddle.float32)
+            place = paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            (result,) = exe.run(fetch_list=[ones])
+            expect = np.ones([10, 2, 3], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
+        with paddle.static.program_guard(paddle.static.Program()):
+            ones = paddle.ones([10, 2, 3], dtype=paddle.float32)
+            place = paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            (result,) = exe.run(fetch_list=[ones])
+            expect = np.ones([10, 2, 3], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
+        with paddle.static.program_guard(paddle.static.Program()):
+            ones = paddle.ones(size=[10, 2, 3], dtype=paddle.float32)
+            place = paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            (result,) = exe.run(fetch_list=[ones])
+            expect = np.ones([10, 2, 3], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
+        with paddle.static.program_guard(paddle.static.Program()):
+            ones = paddle.ones([10, 2, 3], paddle.float32)
+            place = paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            (result,) = exe.run(fetch_list=[ones])
+            expect = np.ones([10, 2, 3], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
+        with paddle.static.program_guard(paddle.static.Program()):
+            ones = paddle.ones([10, 2, 3], paddle.float32)
+            place = paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            (result,) = exe.run(fetch_list=[ones])
+            expect = np.ones([10, 2, 3], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
+        with paddle.static.program_guard(paddle.static.Program()):
+            ones = paddle.ones(shape=[10, 2, 3], dtype=paddle.float32)
+            place = paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            (result,) = exe.run(fetch_list=[ones])
+            expect = np.ones([10, 2, 3], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
         with paddle.static.program_guard(paddle.static.Program()):
             ones = paddle.ones(shape=[10])
             place = paddle.CPUPlace()
             exe = paddle.static.Executor(place)
             (result,) = exe.run(fetch_list=[ones])
-            expected_result = np.ones(10, dtype="float32")
-        self.assertEqual((result == expected_result).all(), True)
+            expect = np.ones(10, dtype="float32")
+        np.testing.assert_equal(result, expect)
 
         with paddle.static.program_guard(paddle.static.Program()):
             ones = paddle.ones(shape=[10], dtype="float64")
             place = paddle.CPUPlace()
             exe = paddle.static.Executor(place)
             (result,) = exe.run(fetch_list=[ones])
-            expected_result = np.ones(10, dtype="float64")
-        self.assertEqual((result == expected_result).all(), True)
+            expect = np.ones(10, dtype="float64")
+        np.testing.assert_equal(result, expect)
 
         with paddle.static.program_guard(paddle.static.Program()):
             ones = paddle.ones(shape=[10], dtype="int64")
             place = paddle.CPUPlace()
             exe = paddle.static.Executor(place)
             (result,) = exe.run(fetch_list=[ones])
-            expected_result = np.ones(10, dtype="int64")
-        self.assertEqual((result == expected_result).all(), True)
+            expect = np.ones(10, dtype="int64")
+        np.testing.assert_equal(result, expect)
 
         with paddle.static.program_guard(paddle.static.Program()):
             ones = paddle.ones(shape=10, dtype="int64")
             place = paddle.CPUPlace()
             exe = paddle.static.Executor(place)
             (result,) = exe.run(fetch_list=[ones])
-            expected_result = np.ones(10, dtype="int64")
-        self.assertEqual((result == expected_result).all(), True)
+            expect = np.ones(10, dtype="int64")
+        np.testing.assert_equal(result, expect)
+        paddle.disable_static()
+
+    def test_dygraph_ones(self):
+        paddle.disable_static()
+        result = paddle.ones(10, dtype=paddle.float32)
+        expect = np.ones([10], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
+        result = paddle.ones(10, 2, 3, dtype=paddle.float32)
+        expect = np.ones([10, 2, 3], dtype="float32")
+        np.testing.assert_equal(result, expect)
+
+        result = paddle.ones([10, 2, 3], dtype=paddle.float32)
+        np.testing.assert_equal(result, expect)
+
+        result = paddle.ones(size=[10, 2, 3], dtype=paddle.float32)
+        np.testing.assert_equal(result, expect)
+
+        result = paddle.ones([10, 2, 3], paddle.float32)
+        np.testing.assert_equal(result, expect)
+
+        result = paddle.ones([10, 2, 3], "float32")
+        np.testing.assert_equal(result, expect)
+
+        result = paddle.ones(shape=[10, 2, 3], dtype=paddle.float32)
+        np.testing.assert_equal(result, expect)
 
 
 if __name__ == "__main__":