From 0a140b44a307c535acd3063656692b482f16ce7b Mon Sep 17 00:00:00 2001 From: viiccwen Date: Fri, 10 Jul 2026 03:46:00 +0000 Subject: [PATCH] [Fix][Relax][ONNX] Import TopK indices as int64 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 --- python/tvm/relax/frontend/onnx/onnx_frontend.py | 4 ++-- tests/python/relax/test_frontend_onnx.py | 10 +++++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index e9b951e0f9f5..be3dd5d4e4f4 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -4413,14 +4413,14 @@ def _impl_v11(cls, bb, inputs, attr, params): 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 f96984e42375..5aab4f558ffd 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(): @@ -6563,7 +6563,9 @@ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0, opset=14 if axes is not None: axes = np.array(axes, dtype=np.int64) node_inputs = ["input", "pads", "", "axes"] - initializer.append(helper.make_tensor("axes", TensorProto.INT64, (len(axes),), axes)) + initializer.append( + helper.make_tensor("axes", TensorProto.INT64, (len(axes),), axes) + ) node = helper.make_node("Pad", inputs=node_inputs, outputs=["output"], mode=mode) graph = helper.make_graph( @@ -6628,7 +6630,9 @@ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0, opset=14 verify_pad( input_shape, pads, - _make_pad_expected_ir(input_shape, pads, mode=mode, value=value, opset=opset, axes=axes), + _make_pad_expected_ir( + input_shape, pads, mode=mode, value=value, opset=opset, axes=axes + ), mode, value, opset,