Skip to content

[rewriter] Unify reshape flatten ops#2518

Merged
justinchuby merged 3 commits intomicrosoft:mainfrom
Johansmm:2301-unify-reshape-flatten-ops
Sep 5, 2025
Merged

[rewriter] Unify reshape flatten ops#2518
justinchuby merged 3 commits intomicrosoft:mainfrom
Johansmm:2301-unify-reshape-flatten-ops

Conversation

@Johansmm
Copy link
Copy Markdown
Contributor

@Johansmm Johansmm commented Aug 26, 2025

Following (#2301), flatten_to_reshape_rule rule set is introduced to reduce the following list of operators:

  • Reshape ∘ Flatten -> Reshape
  • Flatten ∘ Reshape -> Reshape

Note to support this changes:

  • ReshapeReshape rule is updated to support more cases.
  • Flatten2Reshape rule is introduced to convert Flatten ops into Reshape when possible.

@AyoubMDL AyoubMDL mentioned this pull request Aug 27, 2025
Closed
justinchuby
justinchuby reviewed Aug 28, 2025
View reviewed changes
Comment thread onnxscript/rewriter/basic_rules.py
self.new_shape = np.array(np_shape, np_shape.dtype)

# Try to replace {0,-1} values in shape if reshape output is known.
if (reshape_output := context.output_values[0].shape) is not None:
Copy link
Copy Markdown
Collaborator

@justinchuby justinchuby Aug 28, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would this better be done in constant folding? @gramalingam

Or maybe not if a user doesn't want to fold other constants?

Copy link
Copy Markdown
Collaborator

@gramalingam gramalingam Aug 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a different question (and may be that is your main point too): should this rule be decomposed into two simpler rules ... specifically, the transformation applied to the second Reshape in this pattern can be applied even if it does not follow a first Reshape (at least, that is my reading of this, please correct me otherwise). And then we can have the original rule for composing two Reshapes, improved to allow -1 but not zero in the shape.

As to when/where this transformation is done: maybe it could even be done in the torch exporter itself?

@justinchuby justinchuby changed the title 2301 unify reshape flatten ops [rewriter] Unify reshape flatten ops Aug 28, 2025
justinchuby
justinchuby approved these changes Aug 28, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@justinchuby justinchuby left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. @gramalingam and @titaiwangms for more eyes.

justinchuby
justinchuby reviewed Aug 28, 2025
View reviewed changes
Comment thread onnxscript/rewriter/basic_rules.py Outdated
justinchuby
justinchuby reviewed Aug 28, 2025
View reviewed changes
Comment thread onnxscript/rewriter/basic_rules.py Outdated
@codecov
Copy link
Copy Markdown

codecov Bot commented Aug 28, 2025 •
edited
Loading

Codecov Report

❌ Patch coverage is 96.77419% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.16%. Comparing base (1934901) to head (0825709).
⚠️ Report is 4 commits behind head on main.

Files with missing lines Patch % Lines
onnxscript/rewriter/rules/common/_basic_rules.py 94.23% 0 Missing and 3 partials ⚠️
...xscript/rewriter/rules/common/_basic_rules_test.py 98.05% 1 Missing and 1 partial ⚠️
Additional details and impacted files 
@@            Coverage Diff             @@
##             main    #2518      +/-   ##
==========================================
+ Coverage   69.99%   70.16%   +0.16%     
==========================================
  Files         216      216              
  Lines       26074    26215     +141     
  Branches     2618     2638      +20     
==========================================
+ Hits        18250    18393     +143     
+ Misses       6921     6918       -3     
- Partials      903      904       +1

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

xadupre
xadupre reviewed Aug 28, 2025
View reviewed changes
Comment thread onnxscript/rewriter/basic_rules.py Outdated
@Johansmm Johansmm force-pushed the 2301-unify-reshape-flatten-ops branch from 11673eb to 704d4a1 Compare August 28, 2025 20:02
@Johansmm
Copy link
Copy Markdown
Contributor Author

Johansmm commented Aug 28, 2025

Last push force with @justinchuby suggestions

@titaiwangms titaiwangms self-requested a review August 28, 2025 20:04
@Johansmm Johansmm force-pushed the 2301-unify-reshape-flatten-ops branch from 9743c6c to 22bed90 Compare August 28, 2025 21:21
titaiwangms
titaiwangms reviewed Aug 28, 2025
View reviewed changes
Comment thread onnxscript/rewriter/basic_rules.py
@Johansmm Johansmm requested a review from titaiwangms August 29, 2025 18:43
@Johansmm Johansmm force-pushed the 2301-unify-reshape-flatten-ops branch from 22bed90 to 2493299 Compare August 29, 2025 19:53
gramalingam
gramalingam reviewed Aug 30, 2025
View reviewed changes
Comment thread onnxscript/rewriter/basic_rules.py
gramalingam
gramalingam reviewed Aug 30, 2025
View reviewed changes
Comment thread onnxscript/rewriter/basic_rules.py
@Johansmm Johansmm requested a review from gramalingam August 30, 2025 21:47
gramalingam
gramalingam approved these changes Sep 3, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@gramalingam gramalingam left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks fine to me. As discussed in the comments, this could potentially be generalized further by splitting it into two separate rules.

@justinchuby
Copy link
Copy Markdown
Collaborator

justinchuby commented Sep 3, 2025

@Johansmm could you merge from main to resolve conflicts? There was a refactoring done to clean up the name space for all rewrite rules. Thanks

@Johansmm Johansmm force-pushed the 2301-unify-reshape-flatten-ops branch from 144a879 to 2bf14e4 Compare September 4, 2025 07:48
@Johansmm
Copy link
Copy Markdown
Contributor Author

Johansmm commented Sep 4, 2025

Last push force rebasing on main

justinchuby
justinchuby reviewed Sep 4, 2025
View reviewed changes
Comment thread onnxscript/rewriter/rules/common/_basic_rules.py
@Johansmm
[Rewriter] Move reshape rule to another class (microsoft#2301)
653d180 
- rewrite test with ir.tape approach.
- include new tests around check function.
@Johansmm
[Rewrite] Improve ReshapeReshape rule (microsoft#2301)
1c75f3d 
- remove pointless check in shape ignored
- (conditional) support negative shape
- (conditional) support zero shape
@Johansmm
[Rewriter] Introduce Flatten to reshape (microsoft#2301)
0825709 
- Convert Flatten to reshape if possible
- Merge Flatten + Reshape or Reshape + Flatten
@Johansmm Johansmm force-pushed the 2301-unify-reshape-flatten-ops branch from 2bf14e4 to 0825709 Compare September 5, 2025 19:28
@Johansmm
Copy link
Copy Markdown
Contributor Author

Johansmm commented Sep 5, 2025

Last push force exposing rule

@Johansmm Johansmm requested a review from justinchuby September 5, 2025 19:29
@justinchuby justinchuby enabled auto-merge (squash) September 5, 2025 21:17
@justinchuby justinchuby merged commit d0fb218 into microsoft:main Sep 5, 2025
32 checks passed
@github-project-automation github-project-automation Bot moved this from Todo to Done in ONNX Script Review Board Sep 5, 2025
@Johansmm Johansmm deleted the 2301-unify-reshape-flatten-ops branch September 7, 2025 10:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@gramalingam gramalingam gramalingam approved these changes

@justinchuby justinchuby justinchuby approved these changes

@titaiwangms titaiwangms titaiwangms approved these changes

@xadupre xadupre Awaiting requested review from xadupre

Assignees

@gramalingam gramalingam

Labels

module: rewriter

Projects

ONNX Script Review Board
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@Johansmm @justinchuby @gramalingam @titaiwangms @xadupre
⚡