[Optimizer] Collapse Reshape-Transpose-Reshape to Transpose

[Silas Offei Darko (SilasDarko)]

Summary

Implements #1775 by adding a conservative rewrite for:

Reshape -> Transpose -> Reshape

When the intermediate reshape only splits original axes into contiguous groups, the transpose moves those groups without reordering them internally, and the final reshape restores those groups, the sequence is replaced with a single Transpose.

The rule fails closed when those conditions cannot be proven from static shape information.

Scope

This initial implementation supports fully static shapes only.

It intentionally skips chains containing size-0 or size-1 dimensions because those dimensions make axis-boundary provenance ambiguous in the current grouping logic. Symbolic or unknown shapes are also left unchanged.

The rule also validates both the input transpose permutation and the derived permutation before rewriting.

Testing

Added coverage for:

• simple and multi-axis grouped rewrites
• non-trivial and identity permutations
• numerical equivalence
• axis-boundary crossing
• interleaved or internally reordered axis groups
• output shapes that appear valid but have different element ordering
• unknown and symbolic shapes
• size-1 dimensions
• malformed permutations
• integration with MaterializeReshapeShape through optimizer.optimize()

Local validation:

• 22 focused tests passed
• 443 relevant rewriter/optimizer tests passed, with 1 pre-existing skip
• Ruff checks passed
• ONNX checker and numerical-equivalence smoke tests passed

Fixes #1775

[Silas Offei Darko (SilasDarko)]

@microsoft-github-policy-service agree

[Copilot]

Pull request overview

Adds a new conservative optimization rule to the rewriter/optimizer pipeline that collapses Reshape -> Transpose -> Reshape into a single Transpose when (and only when) static shape information proves the reshapes are bookkeeping around a block-preserving transpose, addressing the pattern described in #1775.

Changes:

• Introduces a new rewrite rule that derives a permutation over the original axes by tracking reshape “axis-group” provenance and validating block contiguity + internal order through the transpose.
• Integrates the rule into the default rewriter rule set so it runs in the standard optimizer pipeline.
• Adds extensive unit + integration coverage, including negative/adversarial cases and an end-to-end check through optimizer.optimize().

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
onnxscript/rewriter/rules/common/_reshape_transpose_reshape.py	Implements the provably-correct static-shape-only rewrite from Reshape->Transpose->Reshape to Transpose (or Identity for trivial perms).
onnxscript/rewriter/rules/common/_reshape_transpose_reshape_test.py	Adds targeted positive/negative tests plus an optimizer pipeline integration test to guard correctness and ordering.
onnxscript/rewriter/init.py	Wires the new rule into _DEFAULT_REWRITE_RULES so it participates in optimizer.optimize() flows.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov]

Codecov Report

❌ Patch coverage is 96.15385% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.83%. Comparing base (3ba2bf7) to head (b375da2).

Files with missing lines	Patch %	Lines
...ewriter/rules/common/_reshape_transpose_reshape.py	91.66%	4 Missing and 4 partials ⚠️
...er/rules/common/_reshape_transpose_reshape_test.py	98.78%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3015      +/-   ##
==========================================
+ Coverage   72.64%   72.83%   +0.18%     
==========================================
  Files         265      267       +2     
  Lines       32251    32511     +260     
  Branches     3050     3080      +30     
==========================================
+ Hits        23429    23679     +250     
- Misses       7786     7791       +5     
- Partials     1036     1041       +5     

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