apache / apache/tvm

[Bug] Relax ONNX Min and ArgMin mishandle NaN semantics

Open
#19,668 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 `Min` and `ArgMin` consistently with ONNX Runtime when inputs contain `NaN`.

For `Min`, ONNX Runtime propagates `NaN` when either input element is `NaN`.

For `ArgMin`, TVM should return the same index as ONNX Runtime for inputs containing `NaN`.

### Actual behavior

TVM Relax produces different results from ONNX Runtime:

```
Min ORT: [nan, nan, 4.0, nan]
Min TVM: [7.0, nan, 4.0, nan]

ArgMin ORT: 1
ArgMin TVM: 3
```

For Min, the first element differs: ONNX Runtime returns NaN for min(NaN, 7.0), while TVM returns 7.0.

For ArgMin, ONNX Runtime selects index 1 for the input [8.0, NaN, 3.0, 1.0, 5.0], while TVM returns index 3.

### Environment

TVM: 0.14 environment / Relax ONNX frontend
ONNX Runtime: 1.23
Python: 3.11
Target: llvm
OS: Linux

### Steps to reproduce

```
import numpy as np
import onnx
import onnxruntime as ort
from onnx import TensorProto, helper
import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

def run_tvm(model, feeds):
mod = from_onnx(model, keep_params_in_input=False)

with tvm.transform.PassContext(opt_level=3):
ex = tvm.compile(mod, target=tvm.target.Target("llvm"))

vm = relax.VirtualMachine(ex, tvm.cpu())

out = vm["main"](
*[tvm.runtime.tensor(v, tvm.cpu()) for v in feeds.values()]
)

return (out[0] if isinstance(out, (list, tuple)) else out).numpy()

node = helper.make_node("Min", ["a", "b"], ["y"])

graph = helper.make_graph(
[node],
"g",
[
helper.make_tensor_value_info("a", TensorProto.FLOAT, [4]),
helper.make_tensor_value_info("b", TensorProto.FLOAT, [4]),
],
[helper.make_tensor_value_info("y", TensorProto.FLOAT, [4])],
)

model_min = helper.make_model(
graph,
opset_imports=[helper.make_opsetid("", 17)],
)
model_min.ir_version = 9

a = np.array([np.nan, 12.0, 4.0, np.nan], dtype=np.float32)
b = np.array([7.0, np.nan, 9.0, np.nan], dtype=np.float32)

ort_min = ort.InferenceSession(
model_min.SerializeToString(),
providers=["CPUExecutionProvider"],
).run(None, {"a": a, "b": b})[0]

tvm_min = run_tvm(model_min, {"a": a, "b": b})

print("Min ORT:", ort_min.tolist())
print("Min TVM:", tvm_min.tolist())

node = helper.make_node("ArgMin", ["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_argmin = helper.make_model(
graph,
opset_imports=[helper.make_opsetid("", 17)],
)
model_argmin.ir_version = 9

x = np.array([8.0, np.nan, 3.0, 1.0, 5.0], dtype=np.float32)

ort_argmin = int(
ort.InferenceSession(
model_argmin.SerializeToString(),
providers=["CPUExecutionProvider"],
).run(None, {"x": x})[0]
)

tvm_argmin = int(run_tvm(model_argmin, {"x": x}))

print("ArgMin ORT:", ort_argmin)
print("ArgMin TVM:", tvm_argmin)
```

### Triage

* needs-triage

Contributor guide

No contributing guide indexed for this repository

Research direction

Start at the Relax ONNX frontend entry point, from_onnx, and reproduce the Min and ArgMin cases with the provided Python script. Trace how these operators handle NaN values and compare the compiled TVM results with ONNX Runtime. Done means both operators return the ONNX Runtime results for the shown inputs.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.