This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new a50ab7346f [Fix][Relax][ONNX] Import TopK indices as int64 (#19973)
a50ab7346f is described below
commit a50ab7346f030c4e68f011d3d85df04e0c62b0d3
Author: Vic Wen <[email protected]>
AuthorDate: Sat Jul 11 12:40:19 2026 +0800
[Fix][Relax][ONNX] Import TopK indices as int64 (#19973)
Fixes #19972
ONNX specifies that the second output of TopK, `indices`, has element
type `int64`, and the ONNX TopK operator spec constrains the index
tensor type to `tensor(int64)`:
https://onnx.ai/onnx/operators/onnx__TopK.html
The Relax ONNX frontend previously called `relax.op.topk` without
specifying the output indices dtype, so Relax used its default `int32`
indices.
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 `int64` constants. One example is `TopK -> Div`, where
Relax rejects the binary operation because the imported TopK indices are
`int32` while the divisor is `int64`.
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's
`int64` output.
Verification:
- `uv run --no-sync python -m pytest
tests/python/relax/test_frontend_onnx.py::test_topk -q`
Signed-off-by: viiccwen <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 4 ++--
tests/python/relax/test_frontend_onnx.py | 2 +-
2 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index e9b951e0f9..be3dd5d4e4 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -4413,14 +4413,14 @@ class TopK(OnnxOpConverter):
if sorted != 1:
raise ValueError("TopK sorted must be 1 for Relax frontend")
- return relax.op.topk(data, k, axis, ret_type="both", largest=largest)
+ return relax.op.topk(data, k, axis, ret_type="both", largest=largest,
dtype="int64")
@classmethod
def _impl_v1(cls, bb, inputs, attr, params):
data = inputs[0]
k = attr.get("k", 1)
axis = attr.get("axis", -1)
- return relax.op.topk(data, k, axis, ret_type="both")
+ return relax.op.topk(data, k, axis, ret_type="both", dtype="int64")
class SkipLayerNormalization(OnnxOpConverter):
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index a076f01c4a..5aab4f558f 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -5258,7 +5258,7 @@ def test_topk(axis: int, largest: int):
)
model = helper.make_model(graph, producer_name="topk_test")
- check_correctness(model)
+ check_correctness(model, check_dtypes=True)
def test_expand():
