[PIR] No.43 Migrate paddle.unsqueeze into pir (#57617)

* update unsqueeze in new ir * fix bug * Update test_unsqueeze2_op.py * Update test_unsqueeze_op.py * Update test_unsqueeze2_op.py
PaddlePaddle · Sep 26, 2023 · bd1584b · bd1584b
1 parent d55bb44
commit bd1584b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -2606,7 +2606,7 @@ def unsqueeze(x, axis, name=None):
     """
     input = x
     axes = axis
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         if isinstance(axes, int):
             axes = [axes]
         elif isinstance(axes, Variable):

diff --git a/test/legacy_test/test_unsqueeze2_op.py b/test/legacy_test/test_unsqueeze2_op.py
@@ -43,10 +43,12 @@ def if_enable_cinn(self):
         pass
 
     def test_check_output(self):
-        self.check_output(no_check_set=["XShape"], check_prim=True)
+        self.check_output(
+            no_check_set=["XShape"], check_prim=True, check_new_ir=True
+        )
 
     def test_check_grad(self):
-        self.check_grad(["X"], "Out", check_prim=True)
+        self.check_grad(["X"], "Out", check_prim=True, check_new_ir=True)
 
     def init_test_case(self):
         self.ori_shape = (3, 40)
@@ -135,10 +137,10 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output(no_check_set=["XShape"])
+        self.check_output(no_check_set=["XShape"], check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(["X"], "Out")
+        self.check_grad(["X"], "Out", check_new_ir=True)
 
     def init_test_case(self):
         self.ori_shape = (20, 5)
@@ -196,10 +198,10 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output(no_check_set=["XShape"])
+        self.check_output(no_check_set=["XShape"], check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(["X"], "Out")
+        self.check_grad(["X"], "Out", check_new_ir=True)
 
     def init_test_case(self):
         self.ori_shape = (20, 5)