[PIR] No.38 Migrate paddle.stack into pir (#57477)

liyongchao911 · web-flow · commit c515104d5639 · 2023-09-27T18:59:27.000+08:00
* No.38 Migrate paddle.stack into pir

* update test_stack_op, TestStackBF16Op add new ir check

* update codestyle
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -1844,7 +1844,7 @@ def stack(x, axis=0, name=None):
     """
     axis = 0 if axis is None else axis
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.stack(x, axis)
     else:
         if not isinstance(x, list) and not isinstance(x, tuple):
diff --git a/test/legacy_test/test_stack_op.py b/test/legacy_test/test_stack_op.py
@@ -63,10 +63,12 @@ def setUp(self):
         self.attrs = {'axis': self.axis}
 
     def test_check_output(self):
-        self.check_output(check_prim=True)
+        self.check_output(check_prim=True, check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(self.get_x_names(), 'Y', check_prim=True)
+        self.check_grad(
+            self.get_x_names(), 'Y', check_prim=True, check_new_ir=True
+        )
 
 
 class TestStackOp1(TestStackOpBase):
@@ -187,10 +189,12 @@ def setUp(self):
         self.attrs = {'axis': self.axis}
 
     def test_check_output(self):
-        self.check_output(check_prim=True)
+        self.check_output(check_prim=True, check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(self.get_x_names(), 'Y', check_prim=True)
+        self.check_grad(
+            self.get_x_names(), 'Y', check_prim=True, check_new_ir=True
+        )
 
 
 class TestStackAPIWithLoDTensorArray(unittest.TestCase):
