Summary
When I try to build a minimal TE function that uses te.scan, te.create_prim_func fails with:
TypeError: Unsupported Operation: te.ScanOp. Only te.placeholder and te.compute are allowed for now.
I’m wondering if there’s a plan to support te.ScanOp in te.create_prim_func (and in the new TIR pipeline), and if so, what’s the timeline. My use case is to implement an ONNX GRU/LSTM-like converter in tvm/relax/frontend/onnx similar to the historical TOPI LSTM (which uses te.scan).
Minimal Reproducible Example (TE + te.scan)
# English comments only
import tvm
from tvm import te, tir
def stub_with_scan(T=4, B=2, H=3):
T = tvm.tir.IntImm("int32", T)
B = tvm.tir.IntImm("int32", B)
H = tvm.tir.IntImm("int32", H)
scan_len = tvm.tir.IntImm("int32", int(T) + 1)
# state placeholders
h_state = te.placeholder((scan_len, B, H), name="h_state")
c_state = te.placeholder((scan_len, B, H), name="c_state")
# init (runtime inputs)
h_init = te.placeholder((1, B, H), name="h_init")
c_init = te.placeholder((1, B, H), name="c_init")
def step(prev, init, name):
return te.compute(
(scan_len, B, H),
lambda t, b, j: tir.if_then_else(
t == 0, init[0, b, j], prev[t - 1, b, j]
),
name=name,
)
next_h = step(h_state, h_init, "next_h")
next_c = step(c_state, c_init, "next_c")
scan_h, scan_c = te.scan(
init=[h_init, c_init],
update=[next_h, next_c],
state_placeholder=[h_state, c_state],
name="stub_scan",
inputs=[],
)
# drop t=0 to expose (T, B, H)
hidden = te.compute((T, B, H), lambda t, b, j: scan_h[t + 1, b, j], name="hidden")
cell = te.compute((T, B, H), lambda t, b, j: scan_c[t + 1, b, j], name="cell")
return (hidden, cell), (h_init, c_init)
(hidden, cell), inputs = stub_with_scan()
# This line fails:
prim = te.create_prim_func([*(inputs), hidden, cell])
Observed error
Exception has occurred: TypeError
Unsupported Operation: te.ScanOp. Only te.placeholder and te.compute are allowed for now.
Motivation / Background
- I’m trying to implement
GRUOnnxOpConverterundertvm/relax/frontend/onnxby following the style of historical TOPI LSTM that usedte.scanfor time-step recurrence.
What I tried
- TE without
te.scan: works (builds and runs), which suggests the failure is specifically aboutte.ScanOp. - Port to TIR directly: possible, but I’d like to keep TE parity for GRU/LSTM-like ops (mirroring the older TOPI design), and also avoid prematurely committing to TIR for this converter.
- Relax loop constructs: I’m working in
relax/frontend/onnx, but the TE path is very convenient for expressing RNN-like recurrences ifte.scanwere supported.
Questions
- Is
te.ScanOpplanned to be supported byte.create_prim_funcand the new TIR scheduling pipeline? - If yes, is there a rough timeline or a tracking issue/RFC I can follow?
- If not, what is the recommended path to implement ONNX GRU/LSTM-style recurrences today?
- Should we implement the loop directly in TIR (e.g.,
for/blockwith state carried across time)? - Or is there a Relax-level recommended pattern for dynamic time loops that replaces the older TE
scanapproach?
- For feature parity with legacy TOPI LSTM (which used
te.scan), what is the official guidance in the current architecture?
Why this matters
- ONNX RNN family (LSTM/GRU/RNN) are common; keeping a clean, maintainable way to express their time-step recurrences in TVM helps maintainers and users.
- Many existing examples and mental models for sequence ops were based on TE
scan. Ifscanis no longer a supported path, it would be great to have clear migration guidance (TE→TIR/Relax) and roadmap.