[Bug] Relax ONNX Gather mishandles negative indices
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
TVM Relax should preserve ONNX `Gather` semantics for negative indices.
For ONNX `Gather`, negative indices are valid and should count from the end of the selected axis. For example:
```
X = [1, 2, 3, 4, 5]
I = [-1, -3, 0]
```
The expected output is:
`[5, 3, 1]`
This matches ONNX Runtime.
### Actual behavior
TVM Relax produces a different output after importing the ONNX model and applying LegalizeOps:
```
ORT: [5. 3. 1.]
TVM: [0. 0. 1.]
```
It looks like the negative indices are not handled according to ONNX semantics. In this case, -1 and -3 are effectively producing 0 instead of selecting elements from the end of the input tensor.
### 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
from onnx import helper, TensorProto
import onnxruntime as ort
import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
x_info = helper.make_tensor_value_info("X", TensorProto.FLOAT, [5])
i_info = helper.make_tensor_value_info("I", TensorProto.INT64, [3])
y_info = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3])
node = helper.make_node("Gather", ["X", "I"], ["Y"], axis=0)
graph = helper.make_graph([node], "g", [x_info, i_info], [y_info])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
model.ir_version = 9
onnx.checker.check_model(model)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
idx = np.array([-1, -3, 0], dtype=np.int64)
ort_out = ort.InferenceSession(
model.SerializeToString(),
providers=["CPUExecutionProvider"],
).run(None, {"X": x, "I": idx})[0]
mod = from_onnx(model)
mod = relax.transform.LegalizeOps()(mod)
ex = relax.build(mod, tvm.target.Target("llvm"))
dev = tvm.cpu(0)
vm = relax.VirtualMachine(ex, dev)
tvm_out = vm["main"](
tvm.runtime.tensor(x, device=dev),
tvm.runtime.tensor(idx, device=dev),
).numpy()
print("ORT:", ort_out)
print("TVM:", tvm_out)
```
### Triage
* needs-triage
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the Relax ONNX frontend's from_onnx handling of Gather and the LegalizeOps pass, using the provided Python reproduction to compare ONNX Runtime and TVM. Done means negative indices such as -1 and -3 select elements from the end and the TVM output matches [5., 3., 1.].
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