viiccwen opened a new pull request, #19973: URL: https://github.com/apache/tvm/pull/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: - `TVM_LIBRARY_PATH="$(pwd)/build/lib" uv run --no-sync python -m pytest tests/python/relax/test_frontend_onnx.py::test_topk -q` -- 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]
