Repository navigation
[Bug][Relax][ONNX] Range._impl_v12 rejects runtime Call for limit argument #20064
Description
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).
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.
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.
--- 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. 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.
Activity
zacharywhitley commentedon Jul 29, 2026
Author here. After trying to actually apply the suggested patch, I realized it doesn't work as written — apologies for the misleading sketch. Documenting what I found in case it saves someone else the trip:
bb.emit_te(lambda t: t[()], x)fails at emit time — the lambda returns a Python scalar, not ate.tensor, so it hits"only support te.tensor or tuple/list/Array of te.tensor as function output".te.tensor,bb.emit_tereturns arelax.Var, not aPrimExpr.relax.op.arange's signature isint | PrimExpr | PrimValue, so passing the Var reproduces the original type-mismatch error one hop later.PrimExpr— they live at different phases. The real fix needs either extendingrelax.op.arangeto accept a rank-0 tensor forlimit(with dynamic output shape) or introducing arelax.op.dynamic_arange. Both are substantive changes, not a 3-line patch.I worked around this at the ONNX-graph layer for my use case (fixed-length buckets: rewrite the offending
Rangeto use aConstant(max_frames)forlimit, downstream masking handles positions past the actual duration). Leaving the report open because the underlying issue is real — a large class of transformer ONNX exports hit this pattern — but the fix path is deeper than my initial patch sketch suggested.[Relax][Frontend][ONNX] Support dynamic Range bounds (#20109)