[PIR]Migrate pixel_shuffle to pir (#61096)

changeyoung98 · web-flow · commit 21cd403899a9 · 2024-01-25T14:11:04.000+08:00
* pixel_shuffle to pir passed

* refine
diff --git a/python/paddle/nn/functional/vision.py b/python/paddle/nn/functional/vision.py
@@ -380,7 +380,7 @@ def pixel_shuffle(x, upscale_factor, data_format="NCHW", name=None):
             "Attr(data_format) should be 'NCHW' or 'NHWC'."
             f"But receive Attr(data_format): {data_format} "
         )
-    if in_dygraph_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.pixel_shuffle(x, upscale_factor, data_format)
     else:
         helper = LayerHelper("pixel_shuffle", **locals())
diff --git a/test/legacy_test/test_pixel_shuffle_op.py b/test/legacy_test/test_pixel_shuffle_op.py
@@ -19,8 +19,8 @@
 
 import paddle
 import paddle.nn.functional as F
-from paddle import base
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 def pixel_shuffle_np(x, up_factor, data_format="NCHW"):
@@ -89,12 +89,13 @@ def init_data_format(self):
         self.format = "NCHW"
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
         self.check_grad(
             ['X'],
             'Out',
+            check_pir=True,
         )
 
 
@@ -147,13 +148,14 @@ def init_data_format(self):
         self.format = "NCHW"
 
     def test_check_output(self):
-        self.check_output_with_place(self.place)
+        self.check_output_with_place(self.place, check_pir=True)
 
     def test_check_grad(self):
         self.check_grad_with_place(
             self.place,
             ['X'],
             'Out',
+            check_pir=True,
         )
 
 
@@ -164,6 +166,7 @@ def setUp(self):
         self.out_1_np = pixel_shuffle_np(self.x_1_np, 3)
         self.out_2_np = pixel_shuffle_np(self.x_2_np, 3, "NHWC")
 
+    @test_with_pir_api
     def test_static_graph_functional(self):
         for use_cuda in (
             [False, True] if core.is_compiled_with_cuda() else [False]
@@ -182,22 +185,23 @@ def test_static_graph_functional(self):
 
             exe = paddle.static.Executor(place=place)
             res_1 = exe.run(
-                base.default_main_program(),
-                feed={"x": self.x_1_np},
+                paddle.static.default_main_program(),
+                feed={"x": self.x_1_np, "x2": self.x_2_np},
                 fetch_list=out_1,
                 use_prune=True,
             )[0]
 
             res_2 = exe.run(
-                base.default_main_program(),
-                feed={"x2": self.x_2_np},
+                paddle.static.default_main_program(),
+                feed={"x": self.x_1_np, "x2": self.x_2_np},
                 fetch_list=out_2,
                 use_prune=True,
             )[0]
 
             np.testing.assert_allclose(res_1, self.out_1_np)
             np.testing.assert_allclose(res_2, self.out_2_np)
 
+    @test_with_pir_api
     def test_api_fp16(self):
         paddle.enable_static()
         with paddle.static.program_guard(
@@ -222,21 +226,22 @@ def test_api_fp16(self):
                 out_2_np = pixel_shuffle_np(self.x_2_np, 3, "NHWC")
                 exe = paddle.static.Executor(place=place)
                 res_1 = exe.run(
-                    base.default_main_program(),
-                    feed={"x": self.x_1_np},
+                    paddle.static.default_main_program(),
+                    feed={"x": self.x_1_np, "x2": self.x_2_np},
                     fetch_list=out_1,
                     use_prune=True,
                 )[0]
                 res_2 = exe.run(
-                    base.default_main_program(),
-                    feed={"x2": self.x_2_np},
+                    paddle.static.default_main_program(),
+                    feed={"x": self.x_1_np, "x2": self.x_2_np},
                     fetch_list=out_2,
                     use_prune=True,
                 )[0]
                 np.testing.assert_allclose(res_1, out_1_np)
                 np.testing.assert_allclose(res_2, out_2_np)
 
     # same test between layer and functional in this op.
+    @test_with_pir_api
     def test_static_graph_layer(self):
         for use_cuda in (
             [False, True] if core.is_compiled_with_cuda() else [False]
@@ -260,15 +265,15 @@ def test_static_graph_layer(self):
 
             exe = paddle.static.Executor(place=place)
             res_1 = exe.run(
-                base.default_main_program(),
-                feed={"x": self.x_1_np},
+                paddle.static.default_main_program(),
+                feed={"x": self.x_1_np, "x2": self.x_2_np},
                 fetch_list=out_1,
                 use_prune=True,
             )[0]
 
             res_2 = exe.run(
-                base.default_main_program(),
-                feed={"x2": self.x_2_np},
+                paddle.static.default_main_program(),
+                feed={"x": self.x_1_np, "x2": self.x_2_np},
                 fetch_list=out_2,
                 use_prune=True,
             )[0]

Original file line number	Diff line number	Diff line change
`@@ -380,7 +380,7 @@ def pixel_shuffle(x, upscale_factor, data_format="NCHW", name=None):`
`380`	`380`	`"Attr(data_format) should be 'NCHW' or 'NHWC'."`
`381`	`381`	`f"But receive Attr(data_format): {data_format} "`
`382`	`382`	`)`
`383`		`- if in_dygraph_mode():`
	`383`	`+ if in_dynamic_or_pir_mode():`
`384`	`384`	`return _C_ops.pixel_shuffle(x, upscale_factor, data_format)`
`385`	`385`	`else:`
`386`	`386`	`helper = LayerHelper("pixel_shuffle", **locals())`