[BugFix][Relax]: handle ONNX ScatterElements reduction#19527
[BugFix][Relax]: handle ONNX ScatterElements reduction#19527THINKER-ONLY wants to merge 1 commit intoapache:mainfrom
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for reduction modes in the ScatterElements ONNX operator for opsets 16 and 18 and refactors the reduction attribute processing into a common helper function. Feedback was provided to replace an assert statement with a ValueError to ensure validation is maintained in optimized environments and to improve the error message wording.
| assert reduction in valid_reductions, ( | ||
| f"Only {valid_reductions} reductions are supported, but {reduction} is gotten" | ||
| ) |
There was a problem hiding this comment.
Using assert for validating model attributes is discouraged as it can be optimized away in production environments (when running Python with -O). It is better to raise a ValueError to ensure the validation logic always executes. Additionally, the error message "is gotten" is non-idiomatic; "got" or "received" is preferred.
| assert reduction in valid_reductions, ( | |
| f"Only {valid_reductions} reductions are supported, but {reduction} is gotten" | |
| ) | |
| if reduction not in valid_reductions: | |
| raise ValueError( | |
| f"Only {valid_reductions} reductions are supported, but got {reduction}" | |
| ) |
There was a problem hiding this comment.
Pull request overview
This PR fixes incorrect semantics in the Relax ONNX frontend by honoring the ONNX ScatterElements reduction attribute (introduced in newer opsets) and aligning normalization/validation with existing ScatterND behavior.
Changes:
- Add a shared helper to normalize ONNX
reductionvalues (including mapping ONNX"none"/ missing to Relax"update"). - Implement
ScatterElementsconverters for opset 16 and opset 18 to supportadd/mulandmin/max. - Add regression tests covering multiple opsets and reduction modes for
ScatterElements.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
python/tvm/relax/frontend/onnx/onnx_frontend.py |
Adds reduction normalization helper and implements opset 16/18 ScatterElements conversion; reuses helper for ScatterND. |
tests/python/relax/test_frontend_onnx.py |
Adds parameterized tests validating ScatterElements reduction behavior across opsets. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| assert reduction in valid_reductions, ( | ||
| f"Only {valid_reductions} reductions are supported, but {reduction} is gotten" | ||
| ) |
| @pytest.mark.parametrize( | ||
| "reduction, opset, data, indices, updates", | ||
| [ | ||
| ( | ||
| None, | ||
| 11, | ||
| np.array([[1, 2, 3], [4, 5, 6]], dtype="float32"), | ||
| np.array([[2, 0, 1], [1, 2, 0]], dtype="int64"), | ||
| np.array([[30, 10, 20], [50, 60, 40]], dtype="float32"), | ||
| ), | ||
| ( | ||
| "none", | ||
| 18, | ||
| np.array([[1, 2, 3], [4, 5, 6]], dtype="float32"), | ||
| np.array([[2, 0, 1], [1, 2, 0]], dtype="int64"), | ||
| np.array([[30, 10, 20], [50, 60, 40]], dtype="float32"), | ||
| ), |
THINKER-ONLY
changed the title
Summary
reductionattribute in the Relax ONNX frontendScatterElementsconverter.noneto Relaxupdate.add/mul, and opset 18none/min/max.Changes
ScatterElementsopset 16 and opset 18 converters.relax.op.scatter_elements(..., reduction=...)API.ScatterNDto keep behavior consistent.Test Plan
python -m py_compile python/tvm/relax/frontend/onnx/onnx_frontend.py tests/python/relax/test_frontend_onnx.pypython -m pytest tests/python/relax/test_frontend_onnx.py::test_gather_elements tests/python/relax/test_frontend_onnx.py::test_scatter tests/python/relax/test_frontend_onnx.py::test_scatter_elements_reduction tests/python/relax/test_frontend_onnx.py::test_scatter_nd -qIssue
Fixes #19435
Local Verification Notes
/home/thinker/.cache/tvm-conda-onnx/home/thinker/.cache/tvm-build-onnxtvm.runtime.enabled("llvm") == True14 passed, 4 skipped, 1 warningtests/python/relax/test_frontend_onnx.pywas also attempted. It currently has 14 failures in unrelatedReduce* axes inputandTopKtests; running the same selected failures againstorigin/mainreproduces them, so they are not introduced by this PR.