【PIR API adaptor No.5-6】Migrate paddle.amax/amin into pir (#58546)

GreatV · web-flow · commit 210ca0fa841d · 2023-11-02T14:14:27.000+08:00
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -3079,7 +3079,7 @@ def amax(x, axis=None, keepdim=False, name=None):
              [[0.50000000, 0.33333333],
               [0.        , 0.        ]]])
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.amax(x, axis, keepdim)
 
     else:
@@ -3227,7 +3227,7 @@ def amin(x, axis=None, keepdim=False, name=None):
              [[0.50000000, 0.33333333],
               [0.        , 0.        ]]])
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.amin(x, axis, keepdim)
 
     else:
diff --git a/test/legacy_test/test_max_min_amax_amin_op.py b/test/legacy_test/test_max_min_amax_amin_op.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()
 
@@ -90,6 +91,7 @@ def _choose_paddle_func(self, func, x):
         return out
 
     # We check the output between paddle API and numpy in static graph.
+    @test_with_pir_api
     def test_static_graph(self):
         def _test_static_graph(func):
             startup_program = base.Program()
@@ -103,7 +105,6 @@ def _test_static_graph(func):
 
                 exe = base.Executor(self.place)
                 res = exe.run(
-                    base.default_main_program(),
                     feed={'input': self.x_np},
                     fetch_list=[out],
                 )
