mvanhorn opened a new pull request, #19671: URL: https://github.com/apache/tvm/pull/19671
## Summary ONNX models that run `ArgMax`/`ArgMin` over float tensors containing NaN now import through the Relax frontend with the same winning index that ONNX Runtime returns. ## Why this matters Issue [#19558](https://github.com/apache/tvm/issues/19558) reports that the Relax ONNX frontend picks a different index than ORT and NumPy when NaN is present: ORT lets NaN win `>` comparisons so a leading NaN survives, while `topi.argmax` drops NaN through `val > current_max` and returns the first finite extreme instead. Maintainer swjng confirmed this is an ORT-compatibility gap in the frontend rather than a spec issue, and suggested substituting NaN before the reduction. The fix lives only in the ONNX converter (`onnx_frontend.py`), where ORT parity is the contract, so the generic `topi`/Relax operator is untouched. ## Testing A new `_argreduce_sanitize_nan` helper in `python/tvm/relax/frontend/onnx/onnx_frontend.py` replaces NaN with `+inf` for `ArgMax` and `-inf` for `ArgMin` on floating-point inputs before `argmax`/`argmin`, across the `_impl_v1`/`_impl_v11`/`_impl_v12` paths. Added tests in `tests/python/relax/test_frontend_onnx.py` cover the issue's `[2, NaN, 7, 4, 1]`-style cases against ORT indices `1, 0, 4`, all-NaN and `keepdims=1`, `select_last_index=1`, and a finite-only regression case. Covered by the new tests in this PR; full suite runs in CI. Fixes #19558 -- 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]
