Keep a reducing ScatterND when its static indices cover the whole tensor - #3049
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 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.
ScatterAllStaticrewrites aScatterNDwhose constant indices cover the whole first dimension intoIdentity(updates). Its pattern does not constrainreduction, and unmentioned attributes are unconstrained by default, so the rule also fires onreduction="add" | "mul" | "max" | "min", where the semantics are "combinedatawithupdates" rather than "assign". The rewritten model dropsdataand returnsupdatesalone.ScatterAllDynamic, in the same file, already pinsreduction="none"in its pattern.This is on the default path: the rule is in
_DEFAULT_REWRITE_RULES, so plainoptimize()runs it.torch.onnx.exportdefaults todynamo=True, optimize=True, andx.index_put((idx,), v, accumulate=True)with a constant index exports to exactly this shape, so the whole graph collapses:The existing tests only cover models that omit the attribute, so nothing pinned the reducing case.
The fix is one guard at the top of
check():Tests: four new parameterized cases (
add,mul,max,min) in_redundant_scatter_nd_test.py, each comparing ONNX Runtime output before and afteroptimize(), so the oracle is the unrewritten model rather than my own arithmetic. All four fail onmain.python -m pytest onnxscript/rewriter -qgives 523 passed, 3 skipped.One thing worth flagging: I put the guard in
check()rather than in the pattern whereScatterAllDynamichas it. Writingreduction="none"into the pattern stops it matching aScatterNDthat omits the attribute and relies on the ONNX default (_pattern_ir.py:516-518, withcan_match_noneleft atFalse), which would regress the existingtest_redundant_scatter_nd_static_indices. Happy to move it if you prefer the symmetry and want that test adjusted.