pir support pixel unshuffle op (PaddlePaddle#57521)

jiahy0825 · Oct 16, 2023 · 9daac57 · 9daac57
1 parent 039665a
commit 9daac57
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 137 deletions.
diff --git a/paddle/fluid/operators/pixel_unshuffle_op.cc b/paddle/fluid/operators/pixel_unshuffle_op.cc
diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -1646,6 +1646,15 @@
   kernel :
     func : pixel_shuffle_grad
 
+- backward_op : pixel_unshuffle_grad
+  forward : pixel_unshuffle (Tensor x, int downscale_factor=1, str data_format="NCHW") -> Tensor(out)
+  args : (Tensor out_grad, int downscale_factor, str data_format)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : PixelUnshuffleGradInferMeta
+  kernel :
+    func : pixel_unshuffle_grad
+
 - backward_op : poisson_grad
   forward : poisson (Tensor x) -> Tensor(out)
   args : (Tensor out_grad)

diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -2103,6 +2103,13 @@
   outputs :
     out : Out
 
+- op : pixel_unshuffle
+  backward : pixel_unshuffle_grad
+  inputs :
+    x : X
+  outputs :
+    out : Out
+
 - op : poisson
   inputs :
     x : X

diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -1934,6 +1934,15 @@
     func : pixel_shuffle
   backward : pixel_shuffle_grad
 
+- op : pixel_unshuffle
+  args : (Tensor x, int downscale_factor=1, str data_format="NCHW")
+  output : Tensor
+  infer_meta :
+    func : PixelUnshuffleInferMeta
+  kernel :
+    func : pixel_unshuffle
+  backward : pixel_unshuffle_grad
+
 - op : poisson
   args : (Tensor x)
   output : Tensor

diff --git a/paddle/phi/ops/compat/pixel_unshuffle_sig.cc b/paddle/phi/ops/compat/pixel_unshuffle_sig.cc
diff --git a/test/legacy_test/test_pixel_unshuffle.py b/test/legacy_test/test_pixel_unshuffle.py
@@ -69,8 +69,8 @@ def pixel_unshuffle_np(x, down_factor, data_format="NCHW"):
 
 
 def pixel_unshuffle_wrapper(x, downscale_factor, data_format):
-    return paddle._legacy_C_ops.pixel_unshuffle(
-        x, "downscale_factor", downscale_factor, "data_format", data_format
+    return paddle.nn.functional.pixel_unshuffle(
+        x, downscale_factor, data_format
     )