[API-Compat] More unittest & static graph check & updated decorator

Author: Enigmatisms

python/paddle/tensor/compat.py

@@ -29,12 +29,16 @@
 
     from paddle import Tensor
 
-from paddle.utils.compat_kwarg_check import forbid_keywords
+from paddle.utils.decorator_utils import ForbidKeywordsDecorator
 
 __all__ = []
 
 
-@forbid_keywords(["x", "num_or_sections", "axis", "name"], "paddle.split")
+@ForbidKeywordsDecorator(
+    illegal_keys=["x", "num_or_sections", "axis", "name"],
+    func_name="paddle.compat.split",
+    correct_name="paddle.split",
+)
 def split(
     tensor: Tensor, split_size_or_sections: int | Sequence[int], dim: int = 0
 ) -> tuple[Tensor, ...]:
@@ -105,12 +109,13 @@ def GetSplitSize(split_size, shape_on_dim):
             sections.append(remaining_num)
             return sections
 
-    def SaveGetShapeOnDim(shape, dim: int) -> int:
+    def GetShapeOnDimInRange(shape, dim: int) -> int:
         shape_range = len(shape)
-        if dim < -shape_range or dim >= shape_range:
-            raise ValueError(
-                f"(InvalidArgument) The dim is expected to be in range of [-{shape_range}, {shape_range}), but got {dim}"
-            )
+        if isinstance(dim, int):
+            if dim < -shape_range or dim >= shape_range:
+                raise ValueError(
+                    f"(InvalidArgument) The dim is expected to be in range of [-{shape_range}, {shape_range}), but got {dim}"
+                )
         return shape[dim]
 
     if isinstance(split_size_or_sections, (list, tuple)):
@@ -151,7 +156,7 @@ def SaveGetShapeOnDim(shape, dim: int) -> int:
             ), 'split_size_or_sections must be greater than 0.'
 
             split_size_or_sections = GetSplitSize(
-                split_size_or_sections, SaveGetShapeOnDim(tensor.shape, dim)
+                split_size_or_sections, GetShapeOnDimInRange(tensor.shape, dim)
             )
 
             if isinstance(split_size_or_sections, list):
@@ -164,7 +169,10 @@ def SaveGetShapeOnDim(shape, dim: int) -> int:
             return tuple(_C_ops.split(tensor, split_size_or_sections, dim))
     else:
         if isinstance(dim, paddle.pir.Value):
-            dim.stop_gradient = True
+            raise TypeError(
+                "'dim' is not allowed to be a pir.Value in a static graph: "
+                "\npir.Value can not be used for indexing python lists/tuples."
+            )
         if isinstance(dim, int):
             assert len(tensor.shape) + dim >= 0, "(rank(x) + dim) must >= 0"
             dim = (len(tensor.shape) + dim) if dim < 0 else dim
@@ -173,16 +181,15 @@ def SaveGetShapeOnDim(shape, dim: int) -> int:
 
         if not isinstance(split_size_or_sections, (int, list, tuple)):
             raise TypeError(
-                "The type of 'split_size_or_sections' in split must be int, list or tuple in imperative mode, but "
-                f"received {type(split_size_or_sections)}."
+                "The type of 'split_size_or_sections' in split must be int, list or tuple in imperative mode."
             )
         if isinstance(split_size_or_sections, int):
             assert (
                 split_size_or_sections > 0
             ), 'split_size_or_sections must be greater than 0.'
 
             split_size_or_sections = GetSplitSize(
-                split_size_or_sections, SaveGetShapeOnDim(tensor.shape, dim)
+                split_size_or_sections, GetShapeOnDimInRange(tensor.shape, dim)
             )
             if isinstance(split_size_or_sections, list):
                 if paddle.utils._contain_var(split_size_or_sections):

python/paddle/tensor/manipulation.py

@@ -58,7 +58,7 @@
         TensorOrTensors,
     )
 
-from paddle.utils.compat_kwarg_check import forbid_keywords
+from paddle.utils.decorator_utils import ForbidKeywordsDecorator
 
 __all__ = []
 
@@ -2725,8 +2725,10 @@ def row_stack(x: Sequence[Tensor], name: str | None = None) -> Tensor:
     return paddle.vstack(x, name=name)
 
 
-@forbid_keywords(
-    ["tensor", "split_size_or_sections", "dim"], "paddle.compat.split"
+@ForbidKeywordsDecorator(
+    illegal_keys=["tensor", "split_size_or_sections", "dim"],
+    func_name="paddle.split",
+    correct_name="paddle.compat.split",
 )
 def split(
     x: Tensor,

python/paddle/utils/__init__.py

@@ -14,7 +14,6 @@
 
 from ..base.framework import require_version
 from . import (  # noqa: F401
-    compat_kwarg_check,
     cpp_extension,
     decorator_utils,
     dlpack,

python/paddle/utils/compat_kwarg_check.py

@@ -105,3 +105,32 @@ def process(
                             f"Cannot specify both '{original}' and its alias '{alias}'"
                         )
         return args, processed_kwargs
+
+
+class ForbidKeywordsDecorator(DecoratorBase[_P, _R]):
+    """A decorator that hints users to use the correct `compat` functions, when erroneous keyword arguments are detected"""
+
+    def __init__(
+        self, illegal_keys: list[str] | 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
+
+    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 ""
+
+            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

python/paddle/utils/decorator_utils.py

@@ -71,42 +71,46 @@ def test_chained_operations(self):
         np.testing.assert_allclose(z_np[:, :7], z1.numpy())
         np.testing.assert_allclose(z_np[:, 7:], z2.numpy())
 
-    def test_static_graph(self):
-        """Test static graph execution"""
-        # fixed random seed for reproducibility
-        np.random.seed(114514)
-        # old static graph mode
-        paddle.enable_static()
-
-        with paddle.static.program_guard(paddle.static.Program()):
-            x = paddle.static.data(name='x', shape=[None, 6], dtype='float32')
-            result0, result1 = split(x, split_size_or_sections=[3, 3], dim=1)
-            output = result0 * 2.0 + paddle.sin(result1)
-
-            place = (
-                paddle.CUDAPlace(0)
-                if paddle.is_compiled_with_cuda()
-                else paddle.CPUPlace()
-            )
-            exe = paddle.static.Executor(place)
-
-            input_data = np.random.rand(3, 6).astype('float32')
-            feed = {'x': input_data}
-
-            results = exe.run(feed=feed, fetch_list=[result0, result1, output])
-
-            pd_result0, pd_result1 = results[0], results[1]
-            np.testing.assert_allclose(input_data[:, :3], pd_result0)
-            np.testing.assert_allclose(input_data[:, 3:], pd_result1)
-
-            expected_output = input_data[:, :3] * 2.0 + np.sin(
-                input_data[:, 3:]
-            )
-            np.testing.assert_allclose(
-                expected_output, results[2], rtol=1e-3, atol=1e-3
-            )
-
-        paddle.disable_static()
+    def test_split_grad(self):
+        """Test backprop for split, in1 and in2 are computed by
+        compat.split and original split"""
+
+        def get_tensors():
+            np.random.seed(114514)
+            np_arr = np.random.normal(0, 1, [2, 3, 4, 5])
+            return paddle.to_tensor(np_arr), paddle.to_tensor(np_arr)
+
+        in1, in2 = get_tensors()
+        in1.stop_gradient = False
+        in2.stop_gradient = False
+
+        def computation_graph(in_tensor):
+            y = in_tensor * 2.3 + 3.0
+            y = paddle.maximum(y, paddle.to_tensor([0], dtype=paddle.float32))
+            return y.mean(axis=0)
+
+        out1 = computation_graph(in1)
+        out2 = computation_graph(in2)
+
+        packs1 = paddle.compat.split(out1, 2, dim=2)
+        packs2 = paddle.split(out2, [2, 2, 1], axis=2)
+
+        res1 = packs1[0] + packs1[1] + packs1[2]
+        res2 = packs2[0] + packs2[1] + packs2[2]
+        res1.backward()
+        res2.backward()
+        np.testing.assert_allclose(in1.grad.numpy(), in2.grad.numpy())
+
+    def test_empty_dim(self):
+        """Split with empty dim"""
+        in_tensor = paddle.arange(72, dtype=paddle.int64).reshape([3, 12, 2])
+        self._compare_with_origin(in_tensor, [5, 0, 7], axis=1)
+
+    def test_split_with_one_block(self):
+        """Resulting tuple should be of length 1"""
+        in_tensor = paddle.arange(60, dtype=paddle.float32).reshape([3, 4, 5])
+        self._compare_with_origin(in_tensor, 5, paddle.to_tensor([-1]))
+        self._compare_with_origin(in_tensor, [5], paddle.to_tensor(2))
 
     def test_edge_cases(self):
         """Test edge cases and error handling"""
@@ -131,8 +135,22 @@ def test_error_hint(self):
         """Test whether there will be correct exception when users pass paddle.split kwargs in paddle.compat.split, vice versa."""
         x = paddle.randn([3, 9, 5])
 
-        msg_gt_1 = "split() received unexpected keyword arguments 'tensor', 'split_size_or_sections', 'dim'. \nDid you mean to use paddle.compat.split() instead?"
-        msg_gt_2 = "split() received unexpected keyword argument 'num_or_sections'. \nDid you mean to use paddle.split() instead?"
+        msg_gt_1 = (
+            "paddle.split() received unexpected keyword arguments 'tensor', 'split_size_or_sections', 'dim'. "
+            "\nDid you mean to use paddle.compat.split() instead?"
+        )
+        msg_gt_2 = (
+            "paddle.compat.split() received unexpected keyword argument 'num_or_sections'. "
+            "\nDid you mean to use paddle.split() instead?"
+        )
+        msg_gt_3 = "(InvalidArgument) The dim is expected to be in range of [-3, 3), but got 3"
+        msg_gt_4 = "paddle.compat.split expects split_sizes have only non-negative entries, but got size = -5 on dim 2"
+
+        split_size = paddle.to_tensor([3])
+        msg_gt_5 = (
+            "The type of 'split_size_or_sections' in split must be int, list or tuple in imperative mode, but "
+            f"received {type(split_size)}."
+        )
 
         with self.assertRaises(TypeError) as cm:
             tensors = paddle.split(tensor=x, split_size_or_sections=3, dim=0)
@@ -142,6 +160,18 @@ def test_error_hint(self):
             tensors = split(x, num_or_sections=3, dim=0)
         self.assertEqual(str(cm.exception), msg_gt_2)
 
+        with self.assertRaises(ValueError) as cm:
+            tensors = split(x, 3, dim=3)
+        self.assertEqual(str(cm.exception), msg_gt_3)
+
+        with self.assertRaises(ValueError) as cm:
+            tensors = split(x, [3, 3, -5], -2)
+        self.assertEqual(str(cm.exception), msg_gt_4)
+
+        with self.assertRaises(TypeError) as cm:
+            tensors = split(x, split_size, 1)
+        self.assertEqual(str(cm.exception), msg_gt_5)
+
 
 if __name__ == '__main__':
     unittest.main()