【PIR API adaptor No.258、295、299、307】 Migrate glu/rank/sgn/take into pir (#59535)

DrRyanHuang · web-flow · commit 217cc544bea4 · 2023-12-11T10:49:36.000+08:00
diff --git a/python/paddle/tensor/attribute.py b/python/paddle/tensor/attribute.py
@@ -49,7 +49,7 @@ def rank(input):
             >>> print(rank.numpy())
             3
     """
-    check_type(input, 'input', (Variable), 'input')
+    check_type(input, 'input', (Variable, paddle.pir.Value), 'input')
     ndims = len(input.shape)
     out = assign(np.array(ndims, 'int32'))
 
@@ -163,12 +163,16 @@ def is_complex(x):
             >>> print(paddle.is_complex(x))
             False
     """
-    if not isinstance(x, (paddle.Tensor, paddle.static.Variable)):
+    if not isinstance(
+        x, (paddle.Tensor, paddle.static.Variable, paddle.pir.Value)
+    ):
         raise TypeError(f"Expected Tensor, but received type of x: {type(x)}")
     dtype = x.dtype
     is_complex_dtype = (
         dtype == core.VarDesc.VarType.COMPLEX64
         or dtype == core.VarDesc.VarType.COMPLEX128
+        or dtype == core.DataType.COMPLEX64
+        or dtype == core.DataType.COMPLEX128
     )
     return is_complex_dtype
 
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -6230,6 +6230,11 @@ def sgn(x, name=None):
         paddle.float64,
         paddle.complex64,
         paddle.complex128,
+        DataType.FLOAT16,
+        DataType.FLOAT32,
+        DataType.FLOAT64,
+        DataType.COMPLEX64,
+        DataType.COMPLEX128,
     ]:
         raise TypeError(
             f"The data type of input must be one of ['float16', 'float32', 'float64', 'complex64', 'complex128'], but got {x.dtype}"
@@ -6317,12 +6322,17 @@ def take(x, index, mode='raise', name=None):
             f"'mode' in 'take' should be 'raise', 'wrap', 'clip', but received {mode}."
         )
 
-    if in_dynamic_mode():
-        if not isinstance(index, (paddle.Tensor, Variable)):
+    if in_dynamic_or_pir_mode():
+        if not isinstance(index, (paddle.Tensor, Variable, paddle.pir.Value)):
             raise TypeError(
                 f"The type of 'index' must be Tensor, but got {type(index)}"
             )
-        if index.dtype not in [paddle.int32, paddle.int64]:
+        if index.dtype not in [
+            paddle.int32,
+            paddle.int64,
+            DataType.INT32,
+            DataType.INT64,
+        ]:
             raise TypeError(
                 "The data type of 'index' must be one of ['int32', 'int64'], but got {}".format(
                     index.dtype
diff --git a/test/dygraph_to_static/test_function_spec.py b/test/dygraph_to_static/test_function_spec.py
@@ -106,9 +106,7 @@ def test_args_to_input_spec(self):
         if in_pir_mode():
             self.assertEqual(input_with_spec[1].shape, [4, 10])  # b.shape
         else:
-            self.assertTupleEqual(
-                tuple(input_with_spec[1].shape), (4, 10)
-            )  # b.shape
+            self.assertTupleEqual(input_with_spec[1].shape, (4, 10))  # b.shape
 
         self.assertEqual(input_with_spec[1].name, 'b_var')  # b.name
 
diff --git a/test/legacy_test/test_glu.py b/test/legacy_test/test_glu.py
@@ -20,6 +20,7 @@
 import paddle.base.dygraph as dg
 from paddle import base, nn
 from paddle.nn import functional as F
+from paddle.pir_utils import test_with_pir_api
 
 
 def sigmoid(x):
@@ -58,6 +59,7 @@ def glu_axis_size(self):
         x = paddle.static.data(name='x', shape=[1, 2, 3], dtype='float32')
         paddle.nn.functional.glu(x, axis=256)
 
+    @test_with_pir_api
     def test_errors(self):
         self.assertRaises(ValueError, self.glu_axis_size)
 
@@ -92,6 +94,7 @@ def glu_axis_size(self):
         act = nn.GLU(256)
         act(x)
 
+    @test_with_pir_api
     def test_errors(self):
         self.assertRaises(ValueError, self.glu_axis_size)
         act = nn.GLU(256)
diff --git a/test/legacy_test/test_sgn.py b/test/legacy_test/test_sgn.py
@@ -15,8 +15,10 @@
 import unittest
 
 import numpy as np
+from utils import static_guard
 
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 
 def np_sgn(x: np.ndarray):
@@ -31,7 +33,7 @@ def np_sgn(x: np.ndarray):
 
 
 class TestSgnError(unittest.TestCase):
-    def test_errors(self):
+    def test_errors_dynamic(self):
         # The input dtype of sgn must be float16, float32, float64,complex64,complex128.
         input2 = paddle.to_tensor(
             np.random.randint(-10, 10, size=[12, 20]).astype('int32')
@@ -43,33 +45,28 @@ def test_errors(self):
         self.assertRaises(TypeError, paddle.sgn, input2)
         self.assertRaises(TypeError, paddle.sgn, input3)
 
+    @test_with_pir_api
+    def test_errors_static_and_pir(self):
+        paddle.enable_static()
+        main_program = paddle.static.Program()
+        startup_program = paddle.static.Program()
 
-class TestSignAPI(unittest.TestCase):
-    def setUp(self) -> None:
-        self.support_dtypes = [
-            'float16',
-            'float32',
-            'float64',
-            'complex64',
-            'complex128',
-        ]
-        if paddle.device.get_device() == 'cpu':
-            self.support_dtypes = [
-                'float32',
-                'float64',
-                'complex64',
-                'complex128',
-            ]
-
-    def test_dtype(self):
-        for dtype in self.support_dtypes:
-            x = paddle.to_tensor(
-                np.random.randint(-10, 10, size=[12, 20, 2]).astype(dtype)
+        with paddle.static.program_guard(main_program, startup_program):
+            # The input dtype of sgn must be float16, float32, float64,complex64,complex128.
+            input2 = paddle.to_tensor(
+                np.random.randint(-10, 10, size=[12, 20]).astype('int32')
+            )
+            input3 = paddle.to_tensor(
+                np.random.randint(-10, 10, size=[12, 20]).astype('int64')
             )
 
-            paddle.sgn(x)
+            self.assertRaises(TypeError, paddle.sgn, input2)
+            self.assertRaises(TypeError, paddle.sgn, input3)
+        paddle.disable_static()
+
 
-    def test_complex(self):
+class TestSignAPI(unittest.TestCase):
+    def test_complex_dynamic(self):
         for dtype in ['complex64', 'complex128']:
             np_x = np.array(
                 [[3 + 4j, 7 - 24j, 0, 1 + 2j], [6 + 8j, 3, 0, -2]], dtype=dtype
@@ -80,15 +77,76 @@ def test_complex(self):
             z_expected = np_sgn(np_x)
             np.testing.assert_allclose(np_z, z_expected, rtol=1e-05)
 
-    def test_float(self):
-        for dtype in self.support_dtypes:
+    @test_with_pir_api
+    def test_complex_static_and_pir(self):
+        with static_guard():
+            for dtype in ['complex64', 'complex128']:
+                exe = paddle.static.Executor()
+
+                train_program = paddle.static.Program()
+                startup_program = paddle.static.Program()
+                with paddle.static.program_guard(
+                    train_program, startup_program
+                ):
+                    x = paddle.static.data(name='X', shape=[2, 4], dtype=dtype)
+                    z = paddle.sgn(x)
+
+                # Run the startup program once and only once.
+                # Not need to optimize/compile the startup program.
+                exe.run(startup_program)
+
+                # Run the main program directly without compile.
+                x = np.array(
+                    [[3 + 4j, 7 - 24j, 0, 1 + 2j], [6 + 8j, 3, 0, -2]],
+                    dtype=dtype,
+                )
+                (z,) = exe.run(train_program, feed={"X": x}, fetch_list=[z])
+                z_expected = np_sgn(x)
+                np.testing.assert_allclose(z, z_expected, rtol=1e-05)
+
+    def test_float_dynamic(self):
+        dtype_list = ['float32', 'float64']
+        if paddle.is_compiled_with_cuda():
+            dtype_list.append('float16')
+        for dtype in dtype_list:
             np_x = np.random.randint(-10, 10, size=[12, 20, 2]).astype(dtype)
             x = paddle.to_tensor(np_x)
             z = paddle.sgn(x)
             np_z = z.numpy()
             z_expected = np_sgn(np_x)
             np.testing.assert_allclose(np_z, z_expected, rtol=1e-05)
 
+    @test_with_pir_api
+    def test_float_static_and_pir(self):
+        dtype_list = ['float32', 'float64']
+        if paddle.is_compiled_with_cuda():
+            dtype_list.append('float16')
+        with static_guard():
+            for dtype in dtype_list:
+                exe = paddle.static.Executor()
+
+                train_program = paddle.static.Program()
+                startup_program = paddle.static.Program()
+                with paddle.static.program_guard(
+                    train_program, startup_program
+                ):
+                    np_x = np.random.randint(-10, 10, size=[12, 20, 2]).astype(
+                        dtype
+                    )
+                    x = paddle.static.data(
+                        name='X', shape=[12, 20, 2], dtype=dtype
+                    )
+                    z = paddle.sgn(x)
+
+                # Run the startup program once and only once.
+                # Not need to optimize/compile the startup program.
+                exe.run(startup_program)
+
+                # Run the main program directly without compile.
+                (z,) = exe.run(train_program, feed={"X": np_x}, fetch_list=[z])
+                z_expected = np_sgn(np_x)
+                np.testing.assert_allclose(z, z_expected, rtol=1e-05)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/legacy_test/test_take.py b/test/legacy_test/test_take.py
@@ -18,7 +18,8 @@
 
 import paddle
 from paddle import base
-from paddle.base import Program, core, program_guard
+from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestTakeAPI(unittest.TestCase):
@@ -49,11 +50,12 @@ def setUp(self):
             else base.CPUPlace()
         )
 
+    @test_with_pir_api
     def test_static_graph(self):
         paddle.enable_static()
-        startup_program = Program()
-        train_program = Program()
-        with program_guard(startup_program, train_program):
+        startup_program = paddle.static.Program()
+        train_program = paddle.static.Program()
+        with paddle.static.program_guard(startup_program, train_program):
             x = paddle.static.data(
                 name='input', dtype=self.input_dtype, shape=self.input_shape
             )
@@ -62,9 +64,9 @@ def test_static_graph(self):
             )
             out = paddle.take(x, index, mode=self.mode)
 
-            exe = base.Executor(self.place)
+            exe = paddle.static.Executor(self.place)
             st_result = exe.run(
-                base.default_main_program(),
+                paddle.static.default_main_program(),
                 feed={'input': self.input_np, 'index': self.index_np},
                 fetch_list=out,
             )
@@ -111,10 +113,11 @@ def set_dtype(self):
 class TestTakeTypeError(TestTakeAPI):
     """Test take Type Error"""
 
+    @test_with_pir_api
     def test_static_type_error(self):
         """Argument 'index' must be Tensor"""
         paddle.enable_static()
-        with program_guard(Program()):
+        with paddle.static.program_guard(paddle.static.Program()):
             x = paddle.static.data(
                 name='input', dtype=self.input_dtype, shape=self.input_shape
             )
@@ -127,10 +130,11 @@ def test_dygraph_type_error(self):
         x = paddle.to_tensor(self.input_np)
         self.assertRaises(TypeError, paddle.take, x, self.index_np, self.mode)
 
+    @test_with_pir_api
     def test_static_dtype_error(self):
         """Data type of argument 'index' must be in [paddle.int32, paddle.int64]"""
         paddle.enable_static()
-        with program_guard(Program()):
+        with paddle.static.program_guard(paddle.static.Program()):
             x = paddle.static.data(
                 name='input', dtype='float64', shape=self.input_shape
             )
@@ -178,11 +182,12 @@ def setUp(self):
             else base.CPUPlace()
         )
 
+    @test_with_pir_api
     def test_static_index_error(self):
         """When the index is out of range,
         an error is reported directly through `paddle.index_select`"""
         paddle.enable_static()
-        with program_guard(Program()):
+        with paddle.static.program_guard(paddle.static.Program()):
             x = paddle.static.data(
                 name='input', dtype=self.input_dtype, shape=self.input_shape
             )
