Skip to content

Rewriter: don't turn Flatten into Reshape when a dimension is zero - #3203

Open
dfedoryshchev wants to merge 1 commit into
microsoft:mainfrom
dfedoryshchev:fix/flatten-reshape-zero-dim
Open

dfedoryshchev wants to merge 1 commit into
microsoft:mainfrom
dfedoryshchev:fix/flatten-reshape-zero-dim

Conversation

@dfedoryshchev

Copy link
Copy Markdown
Contributor

Flatten2Reshape turns Flatten into a Reshape whose target shape it computes from the static dims. When one of those dims is 0, the product it writes into the target is 0, and Reshape (without allowzero) reads a 0 as "copy this dim from the input", not as size 0. So a Flatten that works on an empty tensor becomes a Reshape that fails at runtime.
The rule is in basic_optimization_rules() and so in _DEFAULT_REWRITE_RULES, which rewrite(model) and optimize() both run. For example Flatten<axis=1> on float[3,4,0] becomes Reshape(x, [3, 0]), and onnxruntime refuses it with Input shape:{3,4,0}, requested shape:{3,4}; Flatten<axis=3> on float[2,0,3,5] becomes Reshape(x, [0, 5]) and fails the same way (requested shape:{2,5}). The original Flatten runs on both.
The fix is in Flatten2Reshape.check (_basic_rules.py:353): it declines when the computed target has a 0 anywhere except the axis=1 copy of dim 0, which is the one place a copy gives the right size. I went with declining rather than emitting allowzero=1, since the axis=1 path relies on the copy meaning of 0 and allowzero needs opset 14; happy to switch if you prefer that.
test_flatten_to_reshape_with_zero_sized_dim covers both failing shapes plus float[0,3,4] with axis=1, which still rewrites. The first two fail before (the onnxruntime error above) and all three pass after. onnxscript/rewriter, onnxscript/optimizer and onnxscript/version_converter: 884 passed, 6 skipped; ruff format --check clean on the two changed files.

This branch has not been deployed

No deployments
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.

1 participant