Add fuse_consecutive_squeeze_unsqueeze pass#315
Merged
take-cheeze merged 1 commit intoJul 22, 2026
Conversation
PyTorch-to-ONNX export frequently emits a Squeeze immediately followed by an inverse Unsqueeze (or the reverse). When both operators touch exactly the same set of axes they cancel each other out, but no existing pass removed the pair -- fuse_consecutive_squeezes/fuse_consecutive_unsqueezes only handle same-op chains. This adds a Nop pass that eliminates a Squeeze/Unsqueeze pair when their axes match, rewiring the pair's consumers to the original input. Axes are read from either the attribute (opset <= 12) or the input tensor (opset >= 13), and negative axes are normalized against the shared outer rank, recovered from whichever value carries shape information. Fixes onnxsim/onnxsim#315 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Qp39CpZJVZ2wkXuy5Ghwyw Signed-off-by: take-cheeze <takechi101010@gmail.com>
take-cheeze
force-pushed
the
claude/onnx-optimizer-pass-315-cutd99
branch
from
July 22, 2026 05:00
7e5b7a4 to
66e5a3f
Compare
justinchuby
approved these changes
Jul 22, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new ONNX optimizer Nop pass to remove redundant adjacent Squeeze/Unsqueeze (and Unsqueeze/Squeeze) pairs when they operate on the same axes, improving exported-model cleanliness (notably from PyTorch-to-ONNX) and reducing unnecessary ops.
Changes:
- Introduces
fuse_consecutive_squeeze_unsqueezepass to bypass cancelingSqueeze/Unsqueezepairs (including negative-axis normalization). - Registers the new pass in the global pass registry.
- Adds unit tests covering fuse/no-fuse behavior for opset >= 13 (axes-as-input) scenarios.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxoptimizer/test/optimizer_test.py | Adds tests validating elimination of canceling Squeeze/Unsqueeze pairs and ensuring mismatched axes are not fused. |
| onnxoptimizer/passes/fuse_consecutive_squeeze_unsqueeze.h | Implements the new optimization pass that detects and removes canceling adjacent Squeeze/Unsqueeze pairs. |
| onnxoptimizer/pass_registry.h | Wires the new pass into the optimizer registry so it can be invoked by name. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+2887
to
+2894
| def test_fuse_consecutive_squeeze_unsqueeze(self): # type: () -> None | ||
| # Squeeze followed by an inverse Unsqueeze cancels out, leaving only the | ||
| # trailing consumer (Relu) fed directly by the input. | ||
| nodes = [ | ||
| helper.make_node("Squeeze", ["X", "axes"], ["Y"]), | ||
| helper.make_node("Unsqueeze", ["Y", "axes"], ["Z"]), | ||
| helper.make_node("Relu", ["Z"], ["W"]), | ||
| ] |
Comment on lines
+29
to
+31
| #include "onnxoptimizer/pass.h" | ||
| #include "onnxoptimizer/passes/logging.h" | ||
| #include "onnxoptimizer/passes/pass_util.h" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PyTorch-to-ONNX export frequently emits a Squeeze immediately followed by an inverse Unsqueeze (or the reverse). When both operators touch exactly the same set of axes they cancel each other out, but no existing pass removed the pair -- fuse_consecutive_squeezes/fuse_consecutive_unsqueezes only handle same-op chains.
This adds a Nop pass that eliminates a Squeeze/Unsqueeze pair when their axes match, rewiring the pair's consumers to the original input. Axes are read from either the attribute (opset <= 12) or the input tensor (opset
Fixes onnxsim/onnxsim#315
Claude-Session: https://claude.ai/code/session_01Qp39CpZJVZ2wkXuy5Ghwyw