viiccwen opened a new issue, #19972: URL: https://github.com/apache/tvm/issues/19972
### Expected behavior The Relax ONNX frontend should import the second output of ONNX `TopK` as an `int64` tensor. ONNX specifies `TopK` outputs as: - `values`: same element type as the input tensor - `indices`: `int64` The ONNX TopK operator spec constrains the index tensor type to `tensor(int64)`: https://onnx.ai/onnx/operators/onnx__TopK.html This matters when the indices are consumed by later integer or shape/indexing operations such as `Div`, `Gather`, or `Slice`. ### Actual behavior The Relax ONNX frontend calls `relax.op.topk` without specifying the indices dtype. Relax `topk` defaults its indices output to `int32`, so an ONNX graph that uses `TopK` indices together with `int64` constants can fail during import with a Relax type error. For example, a graph containing `TopK -> Div` can fail because the `TopK` indices are imported as `int32`, while the divisor constant is `int64`. ### Environment - Component: Relax ONNX frontend - Affected file: `python/tvm/relax/frontend/onnx/onnx_frontend.py` - Reproduced with a local TVM checkout using the Python ONNX frontend tests ### Steps to reproduce ```python import onnx from onnx import TensorProto, helper from tvm.relax.frontend.onnx import from_onnx k = helper.make_node("Constant", [], ["k"], value=helper.make_tensor("k", TensorProto.INT64, [1], [4])) divisor = helper.make_node( "Constant", [], ["divisor"], value=helper.make_tensor("divisor", TensorProto.INT64, [], [2]) ) topk = helper.make_node("TopK", ["data", "k"], ["values", "indices"], axis=-1) div = helper.make_node("Div", ["indices", "divisor"], ["out"]) graph = helper.make_graph( [k, topk, divisor, div], "topk_indices_int64_test", [helper.make_tensor_value_info("data", TensorProto.FLOAT, [1, 8])], [helper.make_tensor_value_info("out", TensorProto.INT64, [1, 4])], ) model = helper.make_model(graph) from_onnx(model) ``` ### Triage * needs-triage * type: bug * relax * frontend: onnx -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
