[API-Compat] Updated forbid-keyword decorator

Author: Enigmatisms

python/paddle/tensor/compat.py

@@ -16,8 +16,13 @@
 
 from typing import TYPE_CHECKING, NamedTuple
 
+import paddle
 from paddle import _C_ops
 
+from ..framework import (
+    in_dynamic_mode,
+)
+
 if TYPE_CHECKING:
 
     from paddle import Tensor
@@ -32,6 +37,33 @@ class SortRetType(NamedTuple):
     indices: Tensor
 
 
+def _check_out_status(
+    out: Tensor | tuple[Tensor, Tensor] | list[Tensor],
+    expect_multiple: bool = False,
+):
+    if out is None:
+        return
+    if not in_dynamic_mode():
+        raise RuntimeError(
+            "Using `out` static graph CINN backend is currently not supported. Directly return the tensor tuple instead.\n"
+        )
+    if expect_multiple:
+        if not isinstance(out, (tuple, list)) or len(out) != 2:
+            raise TypeError(
+                f"Expected a list or tuple of two tensors, got {type(out)} instead."
+            )
+        if not (
+            isinstance(out[0], paddle.Tensor)
+            and isinstance(out[1], paddle.Tensor)
+        ):
+            raise TypeError(
+                f"Expected Tensor type in the tuple/list, got ({type(out[0])}, {type(out[1])}) instead."
+            )
+    else:
+        if not isinstance(out, paddle.Tensor):
+            raise TypeError(f"Expected a Tensor, got {type(out)} instead.")
+
+
 @ForbidKeywordsDecorator(
     illegal_keys={'x', 'axis'},
     func_name="paddle.compat.sort",
@@ -42,6 +74,7 @@ def sort(
     dim: int = -1,
     descending: bool = False,
     stable: bool = False,
+    out=None,
 ) -> SortRetType:
     """
 
@@ -59,6 +92,8 @@ def sort(
         stable (bool, optional): Whether to use stable sorting algorithm or not.
             When using stable sorting algorithm, the order of equivalent elements
             will be preserved. Default is False.
+        out (tuple, optional) : the output tuple/list of (Tensor, Tensor) that
+            can be optionally given to be used as output buffers
 
     Returns:
         SortRetType, a named tuple which contains `values` and `indices`, can be accessed through either indexing
@@ -112,5 +147,9 @@ def sort(
                      [1, 2, 0, 2],
                      [2, 0, 2, 0]]]))
     """
+    _check_out_status(out, expect_multiple=True)
     outputs, indices = _C_ops.argsort(input, dim, descending, stable)
-    return SortRetType(values=outputs, indices=indices)
+    if out is None:
+        return SortRetType(values=outputs, indices=indices)
+    paddle.assign(outputs, out[0])
+    paddle.assign(indices, out[1])

test/legacy_test/test_compat_sort.py

@@ -17,17 +17,22 @@
 import numpy as np
 
 import paddle
-from paddle.compat import sort
+from paddle.compat import sort as compat_sort
 
 
 class TestCompatSort(unittest.TestCase):
 
     def _compare_with_origin(
-        self, input_tensor, dtype, dim, descending, stable
+        self, input_tensor, dtype, dim, descending, stable, use_out=False
     ):
-        sort_res = sort(
-            input_tensor, dim=dim, descending=descending, stable=stable
-        )
+        """DO NOT set use_out to be True in static graph mode."""
+        if use_out:
+            sort_res = (paddle.to_tensor(0), paddle.to_tensor(0))
+            compat_sort(input_tensor, dim, descending, stable, out=sort_res)
+        else:
+            sort_res = compat_sort(
+                input_tensor, dim=dim, descending=descending, stable=stable
+            )
 
         origin_vals = paddle.sort(
             input_tensor, axis=dim, descending=descending, stable=stable
@@ -37,15 +42,11 @@ def _compare_with_origin(
         )
         if dtype.find("int"):
             np.testing.assert_array_equal(
-                sort_res.values.numpy(), origin_vals.numpy()
+                sort_res[0].numpy(), origin_vals.numpy()
             )
         else:
-            np.testing.assert_allclose(
-                sort_res.values.numpy(), origin_vals.numpy()
-            )
-        np.testing.assert_array_equal(
-            sort_res.indices.numpy(), origin_inds.numpy()
-        )
+            np.testing.assert_allclose(sort_res[0].numpy(), origin_vals.numpy())
+        np.testing.assert_array_equal(sort_res[1].numpy(), origin_inds.numpy())
 
     def test_with_origin_static(self):
         dtypes = [
@@ -75,7 +76,7 @@ def static_graph_tester(descending, stable):
                             input_data = paddle.static.data(
                                 name='x', shape=shape, dtype=dtype
                             )
-                            sort_res = sort(
+                            sort_res = compat_sort(
                                 input_data,
                                 dim=dim,
                                 descending=descending,
@@ -149,19 +150,40 @@ def test_with_origin_dynamic(self, use_static=False):
                     input_tensor = paddle.randint(0, 255, shape).to(dtype)
                 else:
                     input_tensor = paddle.randn(shape, dtype=dtype)
-                for dim in range(len(shape)):
-                    self._compare_with_origin(
-                        input_tensor, dtype, dim, False, False
-                    )
-                    self._compare_with_origin(
-                        input_tensor, dtype, dim - len(shape), False, True
-                    )
-                    self._compare_with_origin(
-                        input_tensor, dtype, dim, True, False
-                    )
-                    self._compare_with_origin(
-                        input_tensor, dtype, dim - len(shape), True, True
-                    )
+                for use_out in [False, True]:
+                    for dim in range(len(shape)):
+                        self._compare_with_origin(
+                            input_tensor,
+                            dtype,
+                            dim,
+                            False,
+                            False,
+                            use_out=use_out,
+                        )
+                        self._compare_with_origin(
+                            input_tensor,
+                            dtype,
+                            dim - len(shape),
+                            False,
+                            True,
+                            use_out=use_out,
+                        )
+                        self._compare_with_origin(
+                            input_tensor,
+                            dtype,
+                            dim,
+                            True,
+                            False,
+                            use_out=use_out,
+                        )
+                        self._compare_with_origin(
+                            input_tensor,
+                            dtype,
+                            dim - len(shape),
+                            True,
+                            True,
+                            use_out=use_out,
+                        )
 
     def test_sort_backward(self):
         """test the backward behavior for all data types"""
@@ -177,7 +199,7 @@ def test_sort_backward(self):
                         y = input_tensor * input_tensor
                     else:
                         y = input_tensor + 1
-                    sort_vals, sort_inds = sort(y, dim=dim)
+                    sort_vals, sort_inds = compat_sort(y, dim=dim)
                     sort_vals.backward()
                     if input_tensor.place.is_gpu_place():
                         np.testing.assert_allclose(
@@ -194,7 +216,7 @@ def test_sort_backward(self):
     def test_edge_cases(self):
         """Test edge cases and error handling"""
         x = paddle.to_tensor([])
-        sort_res = sort(x, descending=True, stable=True)
+        sort_res = compat_sort(x, descending=True, stable=True)
 
         np.testing.assert_array_equal(
             sort_res.values.numpy(), np.array([], dtype=np.float32)
@@ -204,7 +226,7 @@ def test_edge_cases(self):
         )
 
         x = paddle.to_tensor(1)
-        sort_res = sort(input=x, stable=True)
+        sort_res = compat_sort(input=x, stable=True)
 
         np.testing.assert_array_equal(
             sort_res.values.numpy(), np.array(1, dtype=np.float32)
@@ -213,8 +235,8 @@ def test_edge_cases(self):
             sort_res.indices.numpy(), np.array(0, dtype=np.int64)
         )
 
-        msg_gt_1 = "paddle.sort() received unexpected keyword arguments 'input', 'dim'. \nDid you mean to use paddle.compat.sort() instead?"
-        msg_gt_2 = "paddle.compat.sort() received unexpected keyword arguments 'x', 'axis'. \nDid you mean to use paddle.sort() instead?"
+        msg_gt_1 = "paddle.sort() received unexpected keyword arguments 'dim', 'input'. \nDid you mean to use paddle.compat.sort() instead?"
+        msg_gt_2 = "paddle.compat.sort() received unexpected keyword arguments 'axis', 'x'. \nDid you mean to use paddle.sort() instead?"
 
         # invalid split sections
         with self.assertRaises(TypeError) as cm:
@@ -223,9 +245,44 @@ def test_edge_cases(self):
 
         # invalid split axis
         with self.assertRaises(TypeError) as cm:
-            sort(x=paddle.to_tensor([2, 1, 3]), axis=0)
+            compat_sort(x=paddle.to_tensor([2, 1, 3]), axis=0)
         self.assertEqual(str(cm.exception), msg_gt_2)
 
+        def test_wrong_out_input(dim, out_input):
+            with self.assertRaises(TypeError) as cm:
+                compat_sort(paddle.to_tensor([1, 2]), out=out_input)
+
+        test_wrong_out_input(0, [0, paddle.to_tensor(0)])
+        test_wrong_out_input(0, paddle.to_tensor(0))
+        test_wrong_out_input(None, 0)
+        test_wrong_out_input(None, (paddle.to_tensor(0),))
+
+        paddle.enable_static()
+        with (
+            self.assertRaises(RuntimeError) as cm,
+            paddle.static.program_guard(paddle.static.Program()),
+        ):
+            x = paddle.static.data(name='x', shape=[None, 6], dtype='float32')
+            result0, result1 = compat_sort(
+                paddle.arange(24),
+                out=(
+                    paddle.zeros([24]),
+                    paddle.zeros([24], dtype=paddle.int64),
+                ),
+            )
+
+            place = (
+                paddle.CUDAPlace(0)
+                if paddle.is_compiled_with_cuda()
+                else paddle.CPUPlace()
+            )
+            paddle.static.Executor(place).run()
+            self.assertEqual(
+                str(cm.exception),
+                "Using `out` static graph CINN backend is currently not supported. Directly return the tensor tuple instead.\n",
+            )
+        paddle.disable_static()
+
 
 if __name__ == "__main__":
     unittest.main()