Repository navigation
Rewriter: don't turn Flatten into Reshape when a dimension is zero - #3203
Open
dfedoryshchev wants to merge 1 commit into
Open
dfedoryshchev wants to merge 1 commit into
dfedoryshchev wants to merge 1 commit into
Conversation
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Flatten2ReshapeturnsFlatteninto aReshapewhose target shape it computes from the static dims. When one of those dims is 0, the product it writes into the target is 0, andReshape(withoutallowzero) reads a 0 as "copy this dim from the input", not as size 0. So aFlattenthat works on an empty tensor becomes aReshapethat fails at runtime.The rule is in
basic_optimization_rules()and so in_DEFAULT_REWRITE_RULES, whichrewrite(model)andoptimize()both run. For exampleFlatten<axis=1>onfloat[3,4,0]becomesReshape(x, [3, 0]), and onnxruntime refuses it withInput shape:{3,4,0}, requested shape:{3,4};Flatten<axis=3>onfloat[2,0,3,5]becomesReshape(x, [0, 5])and fails the same way (requested shape:{2,5}). The originalFlattenruns on both.The fix is in
Flatten2Reshape.check(_basic_rules.py:353): it declines when the computed target has a 0 anywhere except theaxis=1copy of dim 0, which is the one place a copy gives the right size. I went with declining rather than emittingallowzero=1, since theaxis=1path relies on the copy meaning of 0 andallowzeroneeds opset 14; happy to switch if you prefer that.test_flatten_to_reshape_with_zero_sized_dimcovers both failing shapes plusfloat[0,3,4]withaxis=1, which still rewrites. The first two fail before (the onnxruntime error above) and all three pass after.onnxscript/rewriter,onnxscript/optimizerandonnxscript/version_converter: 884 passed, 6 skipped;ruff format --checkclean on the two changed files.