apache / apache/tvm

[Bug] Relax ONNX ArgMax returns different index from ONNX Runtime when input contains NaN

Open
#19,558 2 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.