Contributed pixelshuffle op (#1514)

Thanks to @justinchuby for a significant amount of help on this :) --------- Co-authored-by: G. Ramalingam <[email protected]> Co-authored-by: Justin Chu <[email protected]>
microsoft · May 14, 2024 · ac7ce49 · ac7ce49
1 parent 1407464
commit ac7ce49
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 3 deletions.
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -6403,10 +6403,23 @@ def aten_pinverse(self: TensorType, rcond: float = 1e-15) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_pixel_shuffle(self: TensorType, upscale_factor: int) -> TensorType:
+@torch_op("aten::pixel_shuffle")
+def aten_pixel_shuffle(self: TReal, upscale_factor: int) -> TReal:
     """pixel_shuffle(Tensor self, int upscale_factor) -> Tensor"""
-
-    raise NotImplementedError()
+    self_shape = op.Shape(self)
+    batch = self_shape[:-3]
+    C_out = op.Unsqueeze(self_shape[-3], [0])
+    H_out = op.Unsqueeze(self_shape[-2], [0])
+    W_out = op.Unsqueeze(self_shape[-1], [0])
+    # Reshaping input by collapsing all leading dimensions to match ONNX op requirement (4D)
+    reshaped_self = op.Reshape(
+        self, op.Concat(op.Unsqueeze(-1, [0]), C_out, H_out, W_out, axis=0)
+    )
+    depth_to_space_output = op.DepthToSpace(
+        reshaped_self, blocksize=upscale_factor, mode="CRD"
+    )
+    output_shape = op.Concat(batch, op.Shape(depth_to_space_output)[1:], axis=0)
+    return op.Reshape(depth_to_space_output, output_shape)
 
 
 def aten_pixel_unshuffle(self: TensorType, downscale_factor: int) -> TensorType:

diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py
@@ -1217,6 +1217,18 @@ def _where_input_wrangler(
         matcher=lambda sample: "weight" in sample.kwargs,
         reason="this Aten overload doesn't accept weight as kwargs",
     ),
+    TorchLibOpInfo(
+        "nn.functional.pixel_shuffle",
+        core_ops.aten_pixel_shuffle,
+    )
+    .xfail(
+        dtypes=(torch.int32, torch.int64),
+        reason="fixme: ONNX Runtime does not support int32/64 inputs",
+    )
+    .xfail(
+        matcher=lambda sample: sample.input.numel() == 0,
+        reason="fixme: ORT does not support empty tensor as input",
+    ),
     TorchLibOpInfo(
         "ops.aten.reflection_pad1d",
         nn_ops.aten_reflection_pad1d,