Fix ReshapeFusion dropping allowzero on inferred 0-sized intermediate dims

[titaiwangms]

Description

ReshapeFusion::FuseContiguousReshapes collapses a chain of Reshape / Squeeze / Unsqueeze nodes into a single Reshape whose shape data is taken verbatim from the fully-inferred output shape of the last node in the chain. The new node is created without an allowzero attribute, so it defaults to allowzero = 0.

When that inferred shape contains a literal 0 dim (legitimate when the original chain used allowzero=1, or when intermediate tensors had zero-sized dimensions), the fused Reshape misinterprets the 0 as "copy the corresponding dim from the input tensor" — but the input here is the original input of the first reshape in the chain, with unrelated dims. The result is a silently wrong output shape (and a benign-looking MergeShapeInfo warning at graph load).

Repro (before the fix)

import numpy as np, onnx, onnxruntime as ort, onnx.reference
from onnx import helper, TensorProto

X  = helper.make_tensor_value_info("X", TensorProto.FLOAT, [0, 6, 2])
Y  = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None, None, None])
s1 = helper.make_tensor("s1", TensorProto.INT64, [3], [3, 2, -1])
s2 = helper.make_tensor("s2", TensorProto.INT64, [3], [0, 0, 3])

n1 = helper.make_node("Reshape", ["X",   "s1"], ["mid"])
n2 = helper.make_node("Reshape", ["mid", "s2"], ["Y"], allowzero=1)
m  = helper.make_model(helper.make_graph([n1, n2], "g", [X], [Y], initializer=[s1, s2]),
                       opset_imports=[helper.make_opsetid("", 18)])

inp = np.random.default_rng(7).random((0, 6, 2), dtype=np.float32)
print("REF:", onnx.reference.ReferenceEvaluator(m).run(None, {"X": inp})[0].shape)
print("ORT:", ort.InferenceSession(m.SerializeToString(),
                                   providers=["CPUExecutionProvider"]).run(None, {"X": inp})[0].shape)

Output on main (40c9f85f69):

REF: (0, 0, 3)
[W ... graph.cc:122 MergeShapeInfo] Error merging shape info for output. 'Y' source:{0,6,3} target:{0,0,3}. Falling back to lenient merge.
ORT: (0, 6, 3)   ❌

Fix

Setting allowzero=1 on the fused node would also work but requires opset >= 14, which this transformer cannot assume (it accepts Reshape opset 5+). Bail out of fusion conservatively when shape_value contains any literal 0 dim.

Test

Adds ReshapeFusionContiguousReshapesWithZeroDim that builds the bug repro programmatically and asserts:

• the two reshapes are NOT collapsed
• the inferred output shape stays (0, 0, 3)

The existing happy-path test ReshapeFusion_Contiguous_Reshape (added in #22494) is unaffected — its inferred output shape (2, 1, 64, 32) contains no zero dims, so the new guard does not trigger.

Provenance

FuseContiguousReshapes was introduced in #22494 (Feb 2025). The bug has been latent in main since then.

Motivation and Context

Found while reviewing microsoft/onnxscript#2907 — the rewriter rule under test there is semantically correct, but its numerical-equivalence check using ORT as the oracle fails because of this fusion bug.

Fixes #28348.

[Copilot]

Pull request overview

This PR fixes a correctness bug in the Reshape fusion optimizer: when collapsing contiguous Reshape/Squeeze/Unsqueeze chains, the fused node could mis-handle literal 0 dimensions because it did not preserve allowzero semantics. In the optimizer stack, this prevents a silent shape corruption during graph rewriting.

Changes:

• Add a guard in ReshapeFusion::FuseContiguousReshapes to skip fusion when the inferred fused shape contains any literal 0 dimension.
• Add a regression test covering a Reshape -> Reshape chain where the second reshape uses allowzero=1 and the correct output shape contains zeros.

Reviewed changes

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

File	Description
onnxruntime/core/optimizer/reshape_fusion.cc	Adds a conservative early-exit to avoid unsafe contiguous reshape fusion for inferred zero dims.
onnxruntime/test/optimizer/graph_transform_test.cc	Adds a regression test asserting the unsafe reshape chain is not fused and inferred output shape stays correct.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[justinchuby]

Will merge to fix the regression reported in #28348. Follow ups can be created if desired.

[justinchuby]

@titaiwangms thanks for creating the fix!