[Bug] Relax ONNX ArgMax returns different index from ONNX Runtime when input contains NaN
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
TVM Relax should execute ONNX `ArgMax` consistently with ONNX Runtime when the input contains `NaN`.
For these inputs, ONNX Runtime matches NumPy's `argmax` behavior:
```
x = [2, NaN, 7, 4, 1] -> index 1
x = [NaN, 2, 7, 4, 1] -> index 0
x = [2, 4, 7, 1, NaN] -> index 4
```
TVM Relax should return the same indices after importing the ONNX model through the Relax ONNX frontend.
### Actual behavior
TVM Relax returns different indices for cases where NaN appears before the finite maximum:
```
> x=[2, NaN, 7, 4, 1] np=1 ORT=1 TVM=2
> x=[NaN, 2, 7, 4, 1] np=0 ORT=0 TVM=2
> x=[2, 4, 7, 1, NaN] np=4 ORT=4 TVM=4
```
The discrepancy appears when importing an ONNX ArgMax model through the Relax ONNX frontend and compiling it for the llvm target.
### Environment
TVM: 0.14 environment / Relax ONNX frontend
ONNX Runtime: 1.23
Python: 3.11
Target: llvm
OS: Linux
### Steps to reproduce
```
import warnings
warnings.filterwarnings("ignore")
import numpy as np
import onnx
import onnxruntime as ort
import tvm
from onnx import TensorProto, helper
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
node = helper.make_node("ArgMax", ["x"], ["y"], axis=0, keepdims=0)
graph = helper.make_graph(
[node],
"g",
[helper.make_tensor_value_info("x", TensorProto.FLOAT, [5])],
[helper.make_tensor_value_info("y", TensorProto.INT64, [])],
)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
model.ir_version = 9
sess = ort.InferenceSession(
model.SerializeToString(),
providers=["CPUExecutionProvider"],
)
mod = from_onnx(model)
with tvm.transform.PassContext(opt_level=3):
ex = tvm.compile(mod, target=tvm.target.Target("llvm"))
vm = relax.VirtualMachine(ex, tvm.cpu())
def run_both(x):
ort_out = int(sess.run(None, {"x": x})[0])
out = vm["main"](tvm.runtime.tensor(x, tvm.cpu()))
tvm_out = int((out[0] if isinstance(out, (list, tuple)) else out).numpy())
numpy_out = int(np.argmax(x))
return numpy_out, ort_out, tvm_out
cases = [
np.array([2.0, np.nan, 7.0, 4.0, 1.0], dtype=np.float32),
np.array([np.nan, 2.0, 7.0, 4.0, 1.0], dtype=np.float32),
np.array([2.0, 4.0, 7.0, 1.0, np.nan], dtype=np.float32),
]
for x in cases:
numpy_out, ort_out, tvm_out = run_both(x)
print(f"x={x.tolist()} np={numpy_out} ORT={ort_out} TVM={tvm_out}")
```
### Triage
* needs-triage
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the provided reproducer and the `from_onnx` entry point, then trace how the ONNX `ArgMax` node is lowered and compiled for the LLVM target. Compare Relax results with NumPy and ONNX Runtime for all three NaN cases; done means the imported model returns the matching indices.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100