Update onnxscript/function_libs/torch_lib/ops/core.py

justinchuby · Copilot · web-flow · commit a7dce4464e90 · 2025-09-03T12:25:11.000-07:00
Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -6726,7 +6726,7 @@ def aten_pixel_unshuffle(self: TReal, downscale_factor: int) -> TReal:
         self, op.Concat(op.Constant(value_ints=[-1]), chw_in_dims, axis=0)
     )
     space_to_depth = op.SpaceToDepth(reshaped_self, blocksize=downscale_factor)
-    final_dims = op.Slice(op.Shape(depth_to_space), starts=[1], ends=[_INT64_MAX], axes=[0])
+    final_dims = op.Slice(op.Shape(space_to_depth), starts=[1], ends=[_INT64_MAX], axes=[0])
     output_shape = op.Concat(batch_dims, final_dims, axis=0)
     return op.Reshape(space_to_depth, output_shape, allowzero=True)
 

Original file line number	Diff line number	Diff line change
`@@ -6726,7 +6726,7 @@ def aten_pixel_unshuffle(self: TReal, downscale_factor: int) -> TReal:`
`6726`	`6726`	`self, op.Concat(op.Constant(value_ints=[-1]), chw_in_dims, axis=0)`
`6727`	`6727`	`)`
`6728`	`6728`	`space_to_depth = op.SpaceToDepth(reshaped_self, blocksize=downscale_factor)`
`6729`		`- final_dims = op.Slice(op.Shape(depth_to_space), starts=[1], ends=[_INT64_MAX], axes=[0])`
	`6729`	`+ final_dims = op.Slice(op.Shape(space_to_depth), starts=[1], ends=[_INT64_MAX], axes=[0])`
`6730`	`6730`	`output_shape = op.Concat(batch_dims, final_dims, axis=0)`
`6731`	`6731`	`return op.Reshape(space_to_depth, output_shape, allowzero=True)`
`6732`	`6732`