[PIR]Migrate trace into pir (#58348)

0x45f · web-flow · commit a7ad947eb4a2 · 2023-10-25T10:51:33.000+08:00
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -3670,7 +3670,7 @@ def __check_input(x, offset, axis1, axis2):
             "But received axis1 = %d, axis2 = %d\n" % (axis1, axis2)
         )
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.trace(x, offset, axis1, axis2)
     else:
         __check_input(x, offset, axis1, axis2)
diff --git a/test/legacy_test/test_trace_op.py b/test/legacy_test/test_trace_op.py
@@ -20,6 +20,7 @@
 import paddle
 from paddle import base, tensor
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestTraceOp(OpTest):
@@ -30,10 +31,10 @@ def setUp(self):
         self.outputs = {'Out': self.target}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['Input'], 'Out')
+        self.check_grad(['Input'], 'Out', check_pir=True)
 
     def init_config(self):
         self.case = np.random.randn(20, 6).astype('float64')
@@ -108,11 +109,15 @@ 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(
-            self.place, ['Input'], 'Out', numeric_grad_delta=0.02
+            self.place,
+            ['Input'],
+            'Out',
+            numeric_grad_delta=0.02,
+            check_pir=True,
         )
 
     def init_config(self):
@@ -145,22 +150,24 @@ def init_config(self):
 
 
 class TestTraceAPICase(unittest.TestCase):
+    @test_with_pir_api
     def test_case1(self):
-        case = np.random.randn(2, 20, 2, 3).astype('float32')
-        data1 = paddle.static.data(
-            name='data1', shape=[2, 20, 2, 3], dtype='float32'
-        )
-        out1 = tensor.trace(data1)
-        out2 = tensor.trace(data1, offset=-5, axis1=1, axis2=-1)
-
-        place = core.CPUPlace()
-        exe = base.Executor(place)
-        results = exe.run(
-            base.default_main_program(),
-            feed={"data1": case},
-            fetch_list=[out1, out2],
-            return_numpy=True,
-        )
+        with paddle.static.program_guard(paddle.static.Program()):
+            case = np.random.randn(2, 20, 2, 3).astype('float32')
+            data1 = paddle.static.data(
+                name='data1', shape=[2, 20, 2, 3], dtype='float32'
+            )
+            out1 = tensor.trace(data1)
+            out2 = tensor.trace(data1, offset=-5, axis1=1, axis2=-1)
+
+            place = core.CPUPlace()
+            exe = base.Executor(place)
+            results = exe.run(
+                paddle.static.default_main_program(),
+                feed={"data1": case},
+                fetch_list=[out1, out2],
+                return_numpy=True,
+            )
         target1 = np.trace(case)
         target2 = np.trace(case, offset=-5, axis1=1, axis2=-1)
         np.testing.assert_allclose(results[0], target1, rtol=1e-05)

Original file line number	Diff line number	Diff line change
`@@ -3670,7 +3670,7 @@ def __check_input(x, offset, axis1, axis2):`
`3670`	`3670`	`"But received axis1 = %d, axis2 = %d\n" % (axis1, axis2)`
`3671`	`3671`	`)`
`3672`	`3672`
`3673`		`- if in_dynamic_mode():`
	`3673`	`+ if in_dynamic_or_pir_mode():`
`3674`	`3674`	`return _C_ops.trace(x, offset, axis1, axis2)`
`3675`	`3675`	`else:`
`3676`	`3676`	`__check_input(x, offset, axis1, axis2)`