Skip to content

fix(optimizer): prevent Shape/Gather from folding to constants when shape has negative (unknown) dims#2964

Open
justinchuby with Copilot wants to merge 2 commits into
mainfrom
copilot/fix-invalid-reshape-error
Open

fix(optimizer): prevent Shape/Gather from folding to constants when shape has negative (unknown) dims#2964
justinchuby with Copilot wants to merge 2 commits into
mainfrom
copilot/fix-invalid-reshape-error

Conversation

Copilot AI commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

FoldConstantsPass incorrectly folds Shape→Gather→Concat→Reshape patterns when the input tensor has unknown dimensions stored as integer -1 in the IR (as opposed to ir.SymbolicDim). Since -1 is an int, the all-integer check passes, causing Shape to fold into Constant(value_ints=[-1, 512, 1, ...]). Gather then returns a constant -1, which propagates into the Reshape shape — producing an invalid shape like [-1, -1] (ONNX allows at most one -1 in a Reshape shape).

Changes

  • Shape handler: add non-negativity guard — only fold to a constant when all shape dims are known positive integers:
    # before
    if all(isinstance(d, int) for d in shape_slice):
    # after
    if all(isinstance(d, int) and d >= 0 for d in shape_slice):
  • Gather handler: same guard on the symbolic-shape gather path, preventing folding when any gathered dimension value is negative.
  • Regression test: test_shape_gather_concat_reshape_with_neg1_dim — verifies that Shape is preserved and Reshape never receives multiple -1 values when the input shape contains -1 integer dims.

The root cause: ONNX IR deserializes unknown dims (dim_value=-1 in the proto) as integer -1, not ir.SymbolicDim. The existing SymbolicDim path was safe; the -1-as-int path was not.

…re negative (-1)

When a tensor shape has `-1` as an integer (representing an unknown/dynamic
dimension), the `Shape` and `Gather` handlers were incorrectly folding to
constants containing `-1`. This caused downstream `Reshape` ops to receive
multiple `-1` values (e.g., `[-1, -1]`), which is invalid in ONNX (at most
one `-1` is allowed in a Reshape shape).

The fix adds a non-negativity check: only fold `Shape`/`Gather` to a constant
when all dimension values are non-negative integers (i.e., all are known static
positive values).

Fixes #2963
Copilot AI changed the title [WIP] Fix invalid reshape issue after FoldConstantsPass fix(optimizer): prevent Shape/Gather from folding to constants when shape has negative (unknown) dims Jul 14, 2026
Copilot AI requested a review from justinchuby July 14, 2026 05:41
@justinchuby justinchuby requested a review from Copilot July 14, 2026 18:09
@justinchuby justinchuby marked this pull request as ready for review July 14, 2026 18:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

This PR hardens the optimizer’s FoldConstantsPass against ONNX-IR shapes that encode unknown dimensions as the integer -1, preventing invalid constant-folding that can propagate multiple -1 values into Reshape shapes.

Changes:

  • Add a non-negativity guard to Shape constant folding so -1 (unknown) dims do not become literal constants.
  • Add the same non-negativity guard to the Gather-from-symbolic-shape folding path to avoid folding negative gathered dims.
  • Add a regression test covering the Shape→Gather→Concat→Reshape pattern when the input shape contains integer -1 dims.

Reviewed changes

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

File Description
onnxscript/optimizer/_constant_folding.py Prevents folding Shape/Gather results to constants when any folded dimension is negative (unknown).
onnxscript/optimizer/_constant_folding_test.py Adds regression coverage for the -1-dim folding bug scenario involving Shape/Gather/Reshape.

Comment on lines +546 to 547
if all(isinstance(d, int) and d >= 0 for d in shape_slice):
return op.Constant(value_ints=ir.AttrInt64s("value_ints", list(shape_slice)))
Comment on lines +709 to +722
model_text = """
<ir_version: 7, opset_import: [ "" : 17]>
agraph (float[N, 512, 4, 1] x) => (float[?] val_5)
{
shape = Shape(x)
index_0 = Constant <value_int=0> ()
gathered = Gather <axis=0> (shape, index_0)
axes = Constant <value_ints=[0]> ()
unsqueezed = Unsqueeze(gathered, axes)
neg_one = Constant <value_ints=[-1]> ()
concat_shape = Concat <axis=0> (unsqueezed, neg_one)
val_5 = Reshape(x, concat_shape)
}
"""
@codecov

codecov Bot commented Jul 14, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 75.00000% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.61%. Comparing base (d3489d4) to head (61b855d).
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
onnxscript/optimizer/_constant_folding_test.py 71.42% 3 Missing and 1 partial ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2964      +/-   ##
==========================================
- Coverage   72.61%   72.61%   -0.01%     
==========================================
  Files         263      263              
  Lines       32034    32048      +14     
  Branches     3013     3014       +1     
==========================================
+ Hits        23263    23273      +10     
- Misses       7748     7751       +3     
- Partials     1023     1024       +1     

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

Invalid reshape after FoldConstantsPass in specific {Shape->Gather->Concat->Reshape} pattern

3 participants