【PIR API adaptor No.7】Migrate paddle.angle into pir (PaddlePaddle#58587)

danleifeng · Nov 2, 2023 · 1cb50fb · 1cb50fb
1 parent 75a5999
commit 1cb50fb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -5943,7 +5943,7 @@ def angle(x, name=None):
              [-1.10714877, -0.78539819,  0.        ,  0.78539819]])
     """
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.angle(x)
     else:
         check_variable_and_dtype(

diff --git a/test/legacy_test/test_angle_op.py b/test/legacy_test/test_angle_op.py
@@ -20,6 +20,7 @@
 import paddle
 from paddle import static
 from paddle.base import core, dygraph
+from paddle.pir_utils import test_with_pir_api
 
 paddle.enable_static()
 
@@ -49,7 +50,7 @@ def setUp(self):
         self.outputs = {'Out': out_ref}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
         self.check_grad(
@@ -58,6 +59,7 @@ def test_check_grad(self):
             user_defined_grads=[
                 angle_grad(self.x, np.ones_like(self.x) / self.x.size)
             ],
+            check_pir=True,
         )
 
 
@@ -93,7 +95,7 @@ def setUp(self):
         self.place = core.CUDAPlace(0)
 
     def test_check_output(self):
-        self.check_output_with_place(self.place)
+        self.check_output_with_place(self.place, check_pir=True)
 
     def test_check_grad(self):
         self.check_grad_with_place(
@@ -103,6 +105,7 @@ def test_check_grad(self):
             user_defined_grads=[
                 angle_grad(self.x, np.ones_like(self.x) / self.x.size)
             ],
+            check_pir=True,
         )
 
 
@@ -119,7 +122,7 @@ def setUp(self):
         self.outputs = {'Out': out_ref}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
         self.check_grad(
@@ -128,6 +131,7 @@ def test_check_grad(self):
             user_defined_grads=[
                 angle_grad(self.x, np.ones_like(self.x) / self.x.size)
             ],
+            check_pir=True,
         )
 
 
@@ -142,6 +146,7 @@ def test_dygraph(self):
             out_np = paddle.angle(x).numpy()
         np.testing.assert_allclose(self.out, out_np, rtol=1e-05)
 
+    @test_with_pir_api
     def test_static(self):
         mp, sp = static.Program(), static.Program()
         with static.program_guard(mp, sp):