fix: align Relax ONNX ArgMax/ArgMin NaN handling with ONNX Runtime

[mvanhorn]

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 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

[gemini-code-assist]

Code Review

This pull request introduces a helper function _argreduce_sanitize_nan to handle NaN values in ONNX ArgMax and ArgMin operators, ensuring they match ONNX Runtime behavior by making NaNs win comparisons. It also adds comprehensive unit tests to verify this behavior under various conditions. The review feedback suggests a valuable optimization and defensive programming improvement: adding a check to prevent potential AttributeError on data.struct_info and implementing compile-time constant folding using NumPy when the input is a relax.Constant.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Optimization & Defensive Programming Opportunity

1. Defensive Check: Added a check if not hasattr(data.struct_info, "dtype") to prevent potential AttributeError if data.struct_info is None or does not have a dtype attribute.
2. Constant Folding: If data is a relax.Constant, we can perform the NaN replacement directly at compile-time using NumPy. This avoids emitting unnecessary isnan, where, and const operators in the Relax graph, keeping the imported IR cleaner and more efficient.

def _argreduce_sanitize_nan(bb, data, *, for_min):
    """Match ONNX Runtime ArgMax/ArgMin behavior by making NaN win comparisons."""
    if not hasattr(data.struct_info, "dtype"):
        return data
    dtype = data.struct_info.dtype
    if not _relax_dtype_is_floating_point(dtype):
        return data
    replacement = -_np.inf if for_min else _np.inf
    if isinstance(data, relax.Constant):
        np_data = data.data.numpy()
        if _np.any(_np.isnan(np_data)):
            np_data = _np.where(_np.isnan(np_data), replacement, np_data)
            return relax.const(np_data, dtype)
        return data
    return bb.emit(
        relax.op.where(relax.op.isnan(data), relax.const(replacement, dtype), data)
    )

[tlopex]

Could we add coverage for NaN together with real infinities? Replacing NaN with +inf/-inf makes it tie with existing infinities, so cases like ArgMax([inf, NaN]) or ArgMin([-inf, NaN]) may no longer match NumPy/ORT tie behavior.

[mvanhorn]

Good idea - added in 28a09fe: test_arg_min_max_nan_with_real_infinities uses rows that mix NaN with real +inf and -inf on the reduction axis, parametrized over ArgMax/ArgMin and both select_last_index values, with correctness checked against ONNX Runtime so the tie-breaking behavior is locked to ort's semantics rather than hand-coded indices.