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

* No.38 Migrate paddle.stack into pir * update test_stack_op, TestStackBF16Op add new ir check * update codestyle
danleifeng · Sep 27, 2023 · 4abf93e · 4abf93e
1 parent ca11326
commit 4abf93e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
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):