Skip to content

[rewriter] Unify reshape flatten ops #2518

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 3 commits into from
Sep 5, 2025
Merged

[rewriter] Unify reshape flatten ops #2518

merged 3 commits into from
Sep 5, 2025

Conversation

Johansmm
Copy link
Contributor

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
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
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
Collaborator

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
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
justinchuby
justinchuby reviewed Aug 28, 2025
View reviewed changes
@codecov Codecov
Copy link

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
@Johansmm Johansmm force-pushed the 2301-unify-reshape-flatten-ops branch from 11673eb to 704d4a1 Compare August 28, 2025 20:02
@Johansmm
Copy link
Contributor Author

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
@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
gramalingam
gramalingam reviewed Aug 30, 2025
View reviewed changes
@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
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
Collaborator

@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
Contributor Author

Johansmm commented Sep 4, 2025

Last push force rebasing on main

justinchuby
justinchuby reviewed Sep 4, 2025
View reviewed changes
@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
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