Repository navigation
JAX: vectorized Scan fails because n_steps reaches the scan as a tracer #2490
Description
Compiling a vectorized Scan with mode="JAX" raises IndexError, because n_steps reaches jax_funcify_Scan as a tracer instead of a Python int. Blockwise.make_node pads the unbatched n_steps into a length-1 MakeVector, and jnp.vectorize hands it to the core scan as an array. Batch size 1 fails the same way, because local_useless_unbatched_blockwise runs after canonicalize and leaves Squeeze(MakeVector(Shape_i)) on n_steps.
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.graph import vectorize_graph
x = pt.vector("x")
a = pt.scalar("a")
ys = pytensor.scan(lambda x_t, a: x_t * a, sequences=[x], non_sequences=[a], return_updates=False)
a_batched = pt.vector("a_batched", shape=(3,))
ys_batched = vectorize_graph(ys, replace={a: a_batched})
fn = pytensor.function([x, a_batched], ys_batched, mode="JAX")
fn(np.arange(5.0), np.ones(3))
# IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax.
# Found slice(None, JitTracer<int64[]>, None).Potentially related to #1425?
Activity
ricardoV94 commentedon Oct 13, 2026
There's a precedent for this in the Blockwise of Convolve (mode symbolic argument). Where we eagerly remove the expand_dims of a critical variable like that so woss backends can take it