Skip to content

Improve Reshape and Slice folding#2807

Open
justinchuby wants to merge 3 commits intomainfrom
justinchu/constant-fold-reshape
Open

Improve Reshape and Slice folding#2807
justinchuby wants to merge 3 commits intomainfrom
justinchu/constant-fold-reshape

Conversation

@justinchuby
Copy link
Collaborator

@justinchuby justinchuby commented Feb 13, 2026

  • Fold the input to reshape if from the output shape we know it can be folded.
  • Eliminate no-op slice if the input shape is the same as output shape.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

This pull request improves constant folding for Reshape operations by materializing shape inputs as constants when the output shape is known from shape inference. It also adds support for INT32 dtype shape inputs, which are common in TFLite models.

Changes:

  • Added INT32 fallback support in get_shape_value() for shape inputs (common in TFLite models)
  • Implemented _try_materialize_reshape_shape() to create constant shape inputs from inferred output shapes
  • Updated reshape() function to leverage shape materialization for optimization

@codecov
Copy link

codecov bot commented Feb 13, 2026

Codecov Report

❌ Patch coverage is 92.48120% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.61%. Comparing base (2b2618e) to head (0bcd3bc).
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
...er/rules/common/_materialize_reshape_shape_test.py 89.61% 1 Missing and 7 partials ⚠️
...nxscript/rewriter/rules/common/_collapse_slices.py 60.00% 1 Missing and 1 partial ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2807      +/-   ##
==========================================
+ Coverage   70.50%   70.61%   +0.11%     
==========================================
  Files         228      230       +2     
  Lines       27114    27246     +132     
  Branches     2726     2746      +20     
==========================================
+ Hits        19116    19240     +124     
- Misses       7065     7066       +1     
- Partials      933      940       +7     

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

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@justinchuby justinchuby force-pushed the justinchu/constant-fold-reshape branch from a67ab69 to 3c550df Compare February 13, 2026 04:37
@justinchuby justinchuby changed the title Improve Reshape constant fold Improve Reshape and Slice folding Feb 13, 2026
@justinchuby justinchuby requested a review from Copilot February 13, 2026 04:38
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Development

Successfully merging this pull request may close these issues.

2 participants