apache / apache/tvm

[Bug] Relax ONNX Gather mishandles negative indices

Open
#19,531 1 comment 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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.