zacharywhitley opened a new issue, #20064:
URL: https://github.com/apache/tvm/issues/20064
## Summary
`Range._impl_v12` in the Relax ONNX frontend passes a `relax.Call` to
`relax.op.arange`'s `limit` parameter, which the FFI signature
rejects. Blocks any ONNX graph with runtime-sized `Range` — attention
masks, position embeddings, dynamic sequence padding. Large class of
transformer models.
## Environment
TVM `0.25.0.post1` (pip), macOS arm64, Python 3.11.
## Reproducer
Model: `owensong/Inflect-Nano-v2` on Hugging Face (Apache-2.0).
```python
import onnx, onnxsim
from tvm.relax.frontend.onnx import from_onnx
model = onnx.load("encoder.onnx")
model, _ = onnxsim.simplify(
model,
overwrite_input_shapes={"tokens": [1, 64], "lengths": [1],
"length_scale": []},
)
mod = from_onnx(model, keep_params_in_input=False)
```
## Failure
```
tvm.error.TVMError: Mismatched type on argument #0 when calling arange:
Expected `ir.PrimExpr` but got `relax.expr.Call`.
```
At
[`python/tvm/relax/frontend/onnx/onnx_frontend.py:3446-3471`](https://github.com/apache/tvm/blob/main/python/tvm/relax/frontend/onnx/onnx_frontend.py).
When Range's `limit` is `ReduceMax(<runtime tensor>)` it arrives as a
`relax.Call`; the frontend hands it to `arange` unchanged; `arange`'s
FFI signature rejects.
## Suggested patch
Lower the runtime `limit` into a scalar `PrimExpr` at the frontend
layer, using `bb.emit_te` to materialize it. Applies to `start` and
`step` symmetrically.
```diff
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ Range._impl_v12
- return relax.op.arange(start, limit, delta, dtype)
+ def _to_prim(x):
+ if isinstance(x, relax.Expr) and not isinstance(x,
(relax.Constant, relax.PrimValue)):
+ # Materialize a rank-0 tensor into a scalar PrimExpr.
+ return bb.emit_te(lambda t: t[()], x)
+ return x
+ return relax.op.arange(_to_prim(start), _to_prim(limit),
_to_prim(delta), dtype)
```
(Alternative: teach `relax.op.arange` itself to accept a rank-0
tensor. More general; benefits every caller. Happy to PR either.)
## No user workaround
Unlike bugs 1 and 2, `onnxsim` doesn't help — `limit` is genuinely a
runtime value derived from tensor contents. Only unblocks are:
patch TVM, or hand-rewrite the ONNX graph to bake in a static max
(fixed-length input buckets).
## Discovered by
Compiling Inflect nano encoder from
[cognition](https://github.com/tegmentum/cognition). Reproducer script
at `crates/inflect-tvm-kernel/scripts/tvm_compile.py`. See sibling
reports for two lower-severity non-structural bugs in the same
frontend that fire on the same model.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
