[Bug][Relax][ONNX] Range._impl_v12 rejects runtime Call for `limit` argument

[zacharywhitley]

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.

[zacharywhitley]

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 a te.tensor, so it hits "only support te.tensor or tuple/list/Array of te.tensor as function output".
• Even with the lambda fixed to return a rank-0 te.tensor, bb.emit_te returns a relax.Var, not a PrimExpr. relax.op.arange's signature is int | PrimExpr | PrimValue, so passing the Var reproduces the original type-mismatch error one hop later.
• Fundamentally, a runtime tensor value can't be lifted into a compile-time PrimExpr — they live at different phases. The real fix needs either extending relax.op.arange to accept a rank-0 tensor for limit (with dynamic output shape) or introducing a relax.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 Range to use a Constant(max_frames) for limit, 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.