Fix Issue #2736

[P-Mahi10]

Fix Issue #2736

This change extends constant folding for Reshape to eliminate dynamic shape
subgraphs of the form Shape → Gather → Concat → Reshape when the target shape
contains exactly one dynamic dimension. The optimizer rewrites the reshape to
use a static -1 dimension, relying on ONNX Reshape semantics to infer the
dynamic value from the input element count. This removes runtime shape
evaluation and enables further canonicalization.

Key Changes

• onnxscript/optimizer/_constant_folding.py
  Updated the reshape optimizer to detect reshape targets with exactly one
  dynamic dimension and replace the dynamic shape computation with a static
  shape containing -1.

• onnxscript/optimizer/_constant_folding_test.py
  Added a unit test to verify that dynamic shape subgraphs (Shape, Gather,
  Concat) are eliminated and that the reshape is either rewritten with -1
  or folded away entirely.

[P-Mahi10]

@microsoft-github-policy-service agree

[codecov]

Codecov Report

❌ Patch coverage is 31.57895% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.41%. Comparing base (3adee71) to head (a101753).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/optimizer/_constant_folding.py	14.28%	16 Missing and 2 partials ⚠️
onnxscript/optimizer/_constant_folding_test.py	52.94%	7 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2769      +/-   ##
==========================================
+ Coverage   70.21%   70.41%   +0.20%     
==========================================
  Files         228      228              
  Lines       27316    27290      -26     
  Branches     2769     2770       +1     
==========================================
+ Hits        19179    19217      +38     
+ Misses       7188     7120      -68     
- Partials      949      953       +4     

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

[Copilot]

Pull request overview

This pull request implements an optimization for the Reshape operator that eliminates dynamic shape computation subgraphs when exactly one dimension is dynamic. The optimization detects patterns like Shape → Gather → Concat → Reshape and replaces them with a static reshape using -1 for the dynamic dimension, relying on ONNX semantics to infer the dimension at runtime.

Changes:

• Extended the reshape optimizer function to detect and optimize single-dynamic-dimension cases
• Added comprehensive test coverage for the new optimization pattern

Reviewed changes

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

File	Description
onnxscript/optimizer/_constant_folding.py	Implements the new reshape optimization logic to replace dynamic shape subgraphs with static -1 dimensions
onnxscript/optimizer/_constant_folding_test.py	Adds test to verify dynamic shape subgraph elimination and correct -1 usage

[Copilot]

The docstring has spelling errors that need to be corrected:

• "f the" should be "If the"
• "imensions" should be "dimensions"
• "untime" should be "runtime"

Suggested change

	f the target shape contains exactly one unknown dimension and all other
	imensions are static, canonicalize the reshape to use -1 and eliminate
	untime shape computation.
	If the target shape contains exactly one unknown dimension and all other
	dimensions are static, canonicalize the reshape to use -1 and eliminate
	runtime shape computation.

[Copilot]

The logic should verify that none of the static dimensions are already -1 before replacing the dynamic dimension with -1. According to ONNX Reshape semantics, at most one dimension can be -1. If the shape already contains a -1 in one of the static positions, this code would create a shape with two -1 dimensions, which is invalid.

Add a check to ensure no static dimension equals -1:

if not isinstance(dim, int) or dim == -1:
    all_others_static = False
    break

Suggested change

	# Verify all other dimensions are concrete integers
	all_others_static = True
	for i, dim in enumerate(shape_value.dims):
	if i not in dynamic_indices:
	if not isinstance(dim, int):
	# Verify all other dimensions are concrete integers and not -1
	all_others_static = True
	for i, dim in enumerate(shape_value.dims):
	if i not in dynamic_indices:
	if not isinstance(dim, int) or dim == -1:

[Copilot]

For consistency with the existing pattern in this file, the Constant creation should wrap the list in ir.AttrInt64s. See lines 452 and 586 in the same file for the established convention:

op.Constant(value_ints=ir.AttrInt64s("value_ints", new_shape_list))

This ensures consistency with how other shape-related constants are created in the optimizer.