[Bug] Relax ONNX Min and ArgMin mishandle NaN semantics
- 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 by tracing the Relax ONNX frontend entry point from_onnx for the Min and ArgMin operators, then run the provided Python reproduction against ONNX Runtime and TVM. Done means Relax matches ONNX Runtime for the shown NaN-containing Min and ArgMin inputs on the LLVM target.
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
- Clearly specified
- Newbie friendliness
- 68/100