apache / apache/tvm

[Bug][Relax][ONNX] BinaryBase.base_impl calls .item() on a TIR PrimExpr

Open Beginner friendly
#20,065 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.