[Fix][Relax][ONNX] Import TopK indices as int64#19973
Open
viiccwen wants to merge 1 commit into
Open
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request updates the ONNX frontend's TopK operator implementations (_impl_v11 and _impl_v1) to explicitly specify dtype="int64" when calling relax.op.topk. Additionally, the corresponding test_topk test is updated to enable data type checking. I have no feedback to provide as there are no review comments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Author
|
cc @guan404ming. 🙌 |
viiccwen
force-pushed
the
fix/onnx-topk-indices-int64
branch
from
80cffed to
319a74b
Compare
ONNX specifies that TopK's second output, indices, has int64 element type. The Relax ONNX frontend previously used relax.op.topk without specifying the indices dtype, so Relax used its default int32 indices. This can make valid ONNX graphs fail or produce an incorrect imported type when TopK indices are consumed by later integer or indexing operations. Pass dtype="int64" when importing ONNX TopK and enable dtype checking in the TopK frontend test. Signed-off-by: viiccwen <vicwen@apache.org>
viiccwen
force-pushed
the
fix/onnx-topk-indices-int64
branch
from
2db4420 to
0a140b4
Compare
guan404ming
approved these changes
Jul 10, 2026
Member
|
LGTM but the ci seems not worked for now. cc @tqchen |
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
3 participants
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.
Fixes #19972
ONNX specifies that the second output of TopK,
indices, has element typeint64, and the ONNX TopK operator spec constrains the index tensor type totensor(int64): https://onnx.ai/onnx/operators/onnx__TopK.htmlThe Relax ONNX frontend previously called
relax.op.topkwithout specifying the output indices dtype, so Relax used its defaultint32indices.This can make otherwise valid ONNX graphs fail during import when the TopK indices are consumed by later integer/index operations that use ONNX's usual
int64constants. One example isTopK -> Div, where Relax rejects the binary operation because the imported TopK indices areint32while the divisor isint64.This patch passes
dtype="int64"when importing ONNX TopK, matching the ONNX operator spec. It also updates the existing TopK frontend test to check output dtypes, so the imported indices must match ONNX Runtime'sint64output.Verification:
uv run --no-sync python -m pytest tests/python/relax/test_frontend_onnx.py::test_topk -q