[torchlib] Improve pixel_shuffle

[justinchuby]

Simplify the graph when input rank is 4, in which case we don't need to do any shape manipulation.

Fix pytorch/pytorch#162061

[Copilot]

Pull Request Overview

This PR optimizes the pixel_shuffle and pixel_unshuffle operations in the torch library by adding fast paths for 4D inputs and improving shape handling for higher-dimensional tensors.

• Adds early return optimization for 4D tensors that can directly use ONNX DepthToSpace/SpaceToDepth operations
• Replaces tensor slicing with explicit op.Slice calls for better shape manipulation
• Adds trace_only=True parameter to both function decorators

[codecov]

Codecov Report

❌ Patch coverage is 28.57143% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.93%. Comparing base (54de741) to head (908ca91).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/core.py	28.57%	8 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2537      +/-   ##
==========================================
- Coverage   69.93%   69.93%   -0.01%     
==========================================
  Files         216      216              
  Lines       26048    26052       +4     
  Branches     2614     2616       +2     
==========================================
+ Hits        18217    18219       +2     
  Misses       6931     6931              
- Partials      900      902       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.