[Hardik K (hardikkaurani)]

Description

This PR improves the ScatterAllDynamic rewrite rule by handling omitted default attributes correctly and validating attribute values before applying the rewrite.

Previously, the pattern explicitly specified Shape(start=0) and ScatterND(reduction="none"), which prevented matching equivalent ONNX nodes where these default attributes were omitted.

The rewrite also needed additional safeguards to ensure that Shape computes the full input dimensions and that ScatterND performs assignment rather than a reduction. Without these checks, applying the rewrite could change model semantics.

Changes

• Default Attribute Matching: Updated the pattern to match Shape and ScatterND without explicitly specifying their default attributes.
• Shape Validation: Added checks to ensure Shape.start is zero and Shape.end is not specified. Attribute references and unsupported values are safely rejected.
• ScatterND Validation: Added checks to ensure the reduction attribute is either omitted or explicitly set to "none". Attribute references and other reduction modes are rejected.
• Regression Tests: Added coverage for successful matching when default attributes are omitted, non-zero Shape.start, the presence of Shape.end, and unsupported ScatterND reduction modes.
• Numerical Validation: Added numerical equivalence testing to verify that the rewritten model preserves the original output.

These changes make the rewrite more robust while preventing transformations when the required semantic conditions are not satisfied.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Hardik K (hardikkaurani)]

Hi maintainers, Justin Chu (@justinchuby) and Ti-Tai Wang (@titaiwangms)

I've opened this PR to improve ScatterAllDynamic matching and make the rewrite safer.

It now handles omitted default attributes for Shape and ScatterND, while explicitly validating the Shape range and ScatterND reduction mode before applying the transformation. I've also added regression tests covering valid default-attribute matching, unsupported attribute configurations, and numerical equivalence.

Thanks for taking a look. I'd appreciate your feedback!

[Copilot]

🟡 Changes recommended

The Shape.start and Shape.end regression tests still pass when the new validation guards are removed because an existing dimension mismatch independently blocks the rewrite.

2 open issues

• Fix fixture so test detects missing Shape.start validation · New
• Fix fixture so test detects missing Shape.end validation · New

_{Review effort: Balanced}

[Hardik K (hardikkaurani)]

Fixed in commit 9eb5bfe6 — updated the fixture to data: float[8, 4], transposed_data: float[8, 4], updates: float[4, 4] with axis = 0 (and added onnx_check(model)), so the existing dimension check passes and Shape.start = 1 is what blocks the rewrite.

[Hardik K (hardikkaurani)]

Fixed in commit 9eb5bfe6 — updated the fixture to data: float[4, 8], transposed_data: float[8, 4], updates: float[4, 4] with axis = -1 (and added onnx_check(model)), so the dimension check passes and the Shape.end guard is directly exercised.

[Hardik K (hardikkaurani)]

Hi Ti-Tai Wang (@titaiwangms), I updated the Shape.start and Shape.end rejection test fixtures in _redundant_scatter_nd_test.py (commit 9eb5bfe6) so that the input/transposed dimension check passes in both cases and the tests directly exercise the new attribute guards (with onnx_check validation added). All _redundant_scatter_nd_test.py and common rewriter tests pass locally. Ready for another look whenever you have time!

[dfedoryshchev]

You missed the same default-attribute case on the Gather: the pattern still has op.Gather(shape, axis, axis=0), and axis=0 becomes an AttrConstantPattern with can_match_none=False, so NodePattern.matches fails with "Attribute axis not found in node" for dim = Gather(shape, axis) with the default axis omitted. The new positive test spells out Gather<axis = 0>, so it does not cover that. Dropping axis=0 from the pattern and checking for an absent or 0 axis in check(), as you did for Shape and ScatterND, would close it.

Minor: Shape<start = -rank> is the full shape too, but start_attr.as_int() != 0 rejects it. That is safe, just stricter than needed.

[Hardik K (hardikkaurani)]

Good catch, dfedoryshchev — thanks for spotting both!

I updated the PR (commits 341bae6b and 58b4e9d7):

• Omitted default Gather.axis: Dropped axis=0 from op.Gather(shape, axis, _outputs=[gather_dim]) in ScatterAllDynamic.pattern and validated in check() that Gather.axis is either absent or 0 (rejecting non-zero values and reference attributes).
• Shape<start = -rank>: Accepted start_val == -len(shape) alongside 0 in ScatterAllDynamic.check() while continuing to reject partial slices (such as start = -1 on a 2D tensor).
• Tests: Updated test_dynamic_indices_with_default_shape_and_scatter_nd_attributes to omit Gather's default axis too, and added tests covering Gather<axis = 1> rejection, Shape<start = -rank> support (with numerical check), and Shape<start = -1> rejection.