Skip to content

Keep a reducing ScatterND when its static indices cover the whole tensor - #3049

Open
dfedoryshchev wants to merge 1 commit into
microsoft:mainfrom
dfedoryshchev:fix-scatter-nd-reduction
Open

dfedoryshchev wants to merge 1 commit into
microsoft:mainfrom
dfedoryshchev:fix-scatter-nd-reduction

Conversation

@dfedoryshchev

Copy link
Copy Markdown
Contributor

ScatterAllStatic rewrites a ScatterND whose constant indices cover the whole first dimension into Identity(updates). Its pattern does not constrain reduction, and unmentioned attributes are unconstrained by default, so the rule also fires on reduction="add" | "mul" | "max" | "min", where the semantics are "combine data with updates" rather than "assign". The rewritten model drops data and returns updates alone. ScatterAllDynamic, in the same file, already pins reduction="none" in its pattern.

This is on the default path: the rule is in _DEFAULT_REWRITE_RULES, so plain optimize() runs it. torch.onnx.export defaults to dynamo=True, optimize=True, and x.index_put((idx,), v, accumulate=True) with a constant index exports to exactly this shape, so the whole graph collapses:

nodes after optimize: ['Identity']
torch          [0,:3]: [1.1011541  0.5783099  0.5431972]
onnx optimized [0,:3]: [0.23550618 0.4039306  0.17603731]

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():

reduction = context.root.attributes.get_string("reduction", "none")
if reduction != "none":
    return result.fail(...)

Tests: four new parameterized cases (add, mul, max, min) in _redundant_scatter_nd_test.py, each comparing ONNX Runtime output before and after optimize(), so the oracle is the unrewritten model rather than my own arithmetic. All four fail on main. python -m pytest onnxscript/rewriter -q gives 523 passed, 3 skipped.

One thing worth flagging: I put the guard in check() rather than in the pattern where ScatterAllDynamic has it. Writing reduction="none" into the pattern stops it matching a ScatterND that omits the attribute and relies on the ONNX default (_pattern_ir.py:516-518, with can_match_none left at False), which would regress the existing test_redundant_scatter_nd_static_indices. Happy to move it if you prefer the symmetry and want that test adjusted.

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