[PIR]Migrate flatten into pir (#57762)

0x45f · web-flow · commit 01f8a7aeac8b · 2023-10-09T10:34:57.000+08:00
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -1549,7 +1549,7 @@ def flatten(x, start_axis=0, stop_axis=-1, name=None):
             img[0, 0, 0, 0] = -1
             print(out[0, 0, 0]) # [-1]
     """
-    if not (isinstance(x, Variable)):
+    if not (isinstance(x, (Variable, paddle.pir.OpResult))):
         raise ValueError("The input x should be a Tensor")
 
     x_dim = len(x.shape)
@@ -1586,7 +1586,7 @@ def flatten(x, start_axis=0, stop_axis=-1, name=None):
         if start_axis > stop_axis:
             raise ValueError("The stop_axis should be larger than stat_axis")
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.flatten(x, start_axis, stop_axis)
     else:
         check_variable_and_dtype(
diff --git a/test/legacy_test/test_flatten_contiguous_range_op.py b/test/legacy_test/test_flatten_contiguous_range_op.py
@@ -46,18 +46,27 @@ def if_enable_cinn(self):
     def test_check_output(self):
         if str(self.dtype) in {"float16", "uint16"}:
             self.check_output_with_place(
-                core.CUDAPlace(0), no_check_set=["XShape"], check_prim=True
+                core.CUDAPlace(0),
+                no_check_set=["XShape"],
+                check_prim=True,
+                check_new_ir=True,
             )
         else:
-            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):
         if str(self.dtype) in {"float16", "uint16"}:
             self.check_grad_with_place(
-                core.CUDAPlace(0), ["X"], "Out", check_prim=True
+                core.CUDAPlace(0),
+                ["X"],
+                "Out",
+                check_prim=True,
+                check_new_ir=True,
             )
         else:
-            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.in_shape = (3, 2, 5, 4)
