Migrate paddle.Tensor.mode into pir (#58883)

ccsuzzh · web-flow · commit ec729e272f6c · 2023-11-10T17:01:59.000+08:00
diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py
@@ -607,7 +607,7 @@ def mode(x, axis=-1, keepdim=False, name=None):
              [2, 1]]))
 
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.mode(x, axis, keepdim)
     else:
         helper = LayerHelper("mode", **locals())
diff --git a/test/legacy_test/test_mode_op.py b/test/legacy_test/test_mode_op.py
@@ -20,6 +20,7 @@
 import paddle
 from paddle import base
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 def _mode1D(a):
@@ -112,12 +113,12 @@ def init_numeric_grads(self):
 
     def test_check_output(self):
         paddle.enable_static()
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
         paddle.enable_static()
         grad = self.init_numeric_grads()
-        self.check_grad({'X'}, 'Out', user_defined_grads=[grad])
+        self.check_grad({'X'}, 'Out', user_defined_grads=[grad], check_pir=True)
 
 
 @unittest.skipIf(
@@ -148,7 +149,7 @@ def test_check_output(self):
         place = core.CUDAPlace(0)
         paddle.enable_static()
         if core.is_bfloat16_supported(place):
-            self.check_output_with_place(place)
+            self.check_output_with_place(place, check_pir=True)
 
     def test_check_grad(self):
         place = core.CUDAPlace(0)
@@ -157,7 +158,7 @@ def test_check_grad(self):
 
         if core.is_bfloat16_supported(place):
             self.check_grad_with_place(
-                place, {'X'}, 'Out', user_defined_grads=[grad]
+                place, {'X'}, 'Out', user_defined_grads=[grad], check_pir=True
             )
 
 
@@ -243,6 +244,7 @@ def setUp(self):
             np.random.random((2, 10, 10)) * 1000, dtype=np.float64
         )
 
+    @test_with_pir_api
     def test_run_static(self):
         paddle.enable_static()
         with paddle.static.program_guard(
