【PIR API adaptor No.156、159、180、189】Migrate some ops into pir (#59041)

zrr1999 · web-flow · commit aac8b6d014a0 · 2023-11-17T10:55:48.000+08:00
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -414,7 +414,7 @@ def multiplex(inputs, index, name=None):
              [3., 4.]])
 
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.multiplex(inputs, index)
     else:
         helper = LayerHelper('multiplex', **locals())
@@ -2406,7 +2406,7 @@ def renorm(x, p, axis, max_norm):
                 )
             )
         axis = axis + len(input_shape)
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         out = _C_ops.renorm(x, p, axis, max_norm)
         return out
     else:
@@ -5420,7 +5420,7 @@ def rad2deg(x, name=None):
             57.29578018)
     """
     rad2deg_scale = 180 / np.pi
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         if convert_dtype(x.dtype) in ['int32', 'int64']:
             x = cast(x, dtype="float32")
         return _C_ops.scale(x, rad2deg_scale, 0.0, True)
@@ -6630,7 +6630,7 @@ def nextafter(x, y, name=None):
             Tensor(shape=[2], dtype=float32, place=Place(cpu), stop_gradient=True,
             [1.00000012, 1.99999988])
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.nextafter(x, y)
     else:
         check_variable_and_dtype(x, 'x', ['float32', 'float64'], 'nextafter')
diff --git a/test/legacy_test/test_multiplex_op.py b/test/legacy_test/test_multiplex_op.py
@@ -45,19 +45,25 @@ def setUp(self):
         self.outputs = {'Out': output}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['x1', 'x2', 'x3', 'x4'], 'Out')
+        self.check_grad(['x1', 'x2', 'x3', 'x4'], 'Out', check_pir=True)
 
     def test_check_grad_ignore_x1(self):
-        self.check_grad(['x2', 'x3', 'x4'], 'Out', no_grad_set=set('x1'))
+        self.check_grad(
+            ['x2', 'x3', 'x4'], 'Out', no_grad_set=set('x1'), check_pir=True
+        )
 
     def test_check_grad_ignore_x1_x2(self):
-        self.check_grad(['x3', 'x4'], 'Out', no_grad_set={'x1', 'x2'})
+        self.check_grad(
+            ['x3', 'x4'], 'Out', no_grad_set={'x1', 'x2'}, check_pir=True
+        )
 
     def test_check_grad_ignore_x3(self):
-        self.check_grad(['x1', 'x2', 'x4'], 'Out', no_grad_set=set('x3'))
+        self.check_grad(
+            ['x1', 'x2', 'x4'], 'Out', no_grad_set=set('x3'), check_pir=True
+        )
 
 
 class TestMultiplexOpError(unittest.TestCase):
diff --git a/test/legacy_test/test_nextafter_op.py b/test/legacy_test/test_nextafter_op.py
@@ -18,6 +18,7 @@
 from op_test import OpTest
 
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 
 def ref_nextafter(x, y):
@@ -39,6 +40,7 @@ def setUp(self):
             else paddle.CPUPlace()
         )
 
+    @test_with_pir_api
     def test_static_api(self):
         paddle.enable_static()
         with paddle.static.program_guard(paddle.static.Program()):
@@ -103,7 +105,7 @@ def setUp(self):
         self.outputs = {'out': out}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def init_dtype(self):
         self.dtype = np.float64
diff --git a/test/legacy_test/test_rad2deg.py b/test/legacy_test/test_rad2deg.py
@@ -19,6 +19,7 @@
 import paddle
 from paddle import base
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 paddle.enable_static()
 
@@ -32,10 +33,11 @@ def setUp(self):
         self.x_shape = [6]
         self.out_np = np.rad2deg(self.x_np)
 
+    @test_with_pir_api
     def test_static_graph(self):
-        startup_program = base.Program()
-        train_program = base.Program()
-        with base.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.x_dtype, shape=self.x_shape
             )
@@ -48,11 +50,10 @@ def test_static_graph(self):
             )
             exe = base.Executor(place)
             res = exe.run(
-                base.default_main_program(),
                 feed={'input': self.x_np},
                 fetch_list=[out],
             )
-            self.assertTrue((np.array(out[0]) == self.out_np).all())
+            np.testing.assert_allclose(self.out_np, res[0], rtol=1e-05)
 
     def test_dygraph(self):
         paddle.disable_static()
@@ -96,3 +97,7 @@ def test_dygraph(self):
         np.testing.assert_allclose(180 / np.pi, result2.numpy(), rtol=1e-05)
 
         paddle.enable_static()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/legacy_test/test_renorm_op.py b/test/legacy_test/test_renorm_op.py
@@ -18,7 +18,7 @@
 
 import paddle
 from paddle import base
-from paddle.base import Program, program_guard
+from paddle.pir_utils import test_with_pir_api
 
 paddle.set_device('cpu')
 
@@ -32,12 +32,15 @@ def input_data(self):
         self.dim = 2
         self.max_norm = 2.05
 
+    @test_with_pir_api
     def test_renorm_api(self):
         paddle.enable_static()
         self.input_data()
 
         # case 1:
-        with program_guard(Program(), Program()):
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
             x = paddle.static.data(name="x", shape=[-1, 2, 3], dtype='float64')
             z = paddle.renorm(x, self.p, self.dim, self.max_norm)
             exe = base.Executor(base.CPUPlace())
