[Bug][Relax][ONNX] BinaryBase.base_impl calls .item() on a TIR PrimExpr
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
## Summary
`BinaryBase.base_impl` in the Relax ONNX frontend crashes when both
operands are `PrimValue`s from shape arithmetic. The numpy fallback
returns a TIR `PrimExpr`, then `.item()` is called on it.
## Environment
TVM `0.25.0.post1` (pip), macOS arm64, Python 3.11.
## Reproducer
Model: `owensong/Inflect-Nano-v2` on Hugging Face (Apache-2.0). The
original (un-simplified) encoder has a `Sub` on two shape-derived
`PrimValue`s.
```python
import onnx
from tvm.relax.frontend.onnx import from_onnx
model = onnx.load("encoder.onnx") # do NOT run onnxsim
mod = from_onnx(model, keep_params_in_input=False)
```
## Failure
```
AttributeError: 'Sub' object has no attribute 'item'
```
At [`python/tvm/relax/frontend/onnx/onnx_frontend.py:475-493`](https://github.com/apache/tvm/blob/main/python/tvm/relax/frontend/onnx/onnx_frontend.py).
`_to_numpy` wraps two `PrimValue`s as 0-d object arrays; numpy dispatch
returns a bare `PrimExpr` (symbolic, not numeric); fall-through calls
`.item()` which the `PrimExpr` doesn't implement.
## Suggested patch
Detect when the numpy op returned a `PrimExpr` and wrap it back into
a `PrimValue` rather than calling `.item()`.
```diff
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ BinaryBase.base_impl
- return output.item()
+ if isinstance(output, tvm.tir.PrimExpr):
+ return relax.PrimValue(output)
+ return output.item()
```
## User workaround
`onnxsim` with concrete input shapes constant-folds most `Shape → Gather
→ Sub` patterns away. Works when input shapes can be pinned; not
viable for genuinely dynamic dimensions.
## Discovered by
Compiling Inflect nano encoder from
[cognition](https://github.com/tegmentum/cognition). See sibling
reports for related structural bugs in the same frontend (bug 3 in
particular has no user workaround).
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in python/tvm/relax/frontend/onnx/onnx_frontend.py at BinaryBase.base_impl, lines 475-493, and run the supplied from_onnx reproducer with the unsimplified encoder. Trace the numpy fallback for the two PrimValue operands; done means the dynamic Shape → Gather → Sub path no longer raises AttributeError when it returns a TIR PrimExpr.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100