Follow-up in the same area as #134 and #193/#196. The WhileLoopPrimitive from #134 fixed the graph-growth crash, but a scan that stacks outputs — which is how JAX collects MCMC samples, RNN hidden states, filter histories — now pays a per-iteration cost that grows with T, so total time is quadratic. A scan that carries the same state without stacking is flat.
MWE:
import time
import jax
import jax.numpy as jnp
def bench(f, *args):
o = f(*args); jax.block_until_ready(o)
t0 = time.perf_counter()
o = f(*args); jax.block_until_ready(o)
return time.perf_counter() - t0
x0 = jnp.linspace(-1.0, 1.0, 1_000_000)
for T in (25, 50, 100, 200):
def no_stack(x, T=T):
def step(c, _):
return jnp.tanh(0.99 * c + 0.01), None
c, _ = jax.lax.scan(step, x, None, length=T)
return c
def stack(x, T=T):
def step(c, _):
c2 = jnp.tanh(0.99 * c + 0.01)
return c2, c2
_, ys = jax.lax.scan(step, x, None, length=T)
return ys[-1]
a = bench(jax.jit(no_stack), x0)
b = bench(jax.jit(stack), x0)
print(f"T={T:4d} no-stack {a*1e3:7.1f} ms stacked {b*1e3:7.1f} ms"
f" stack cost/step {(b-a)/T*1e3:6.2f} ms")Results (M3 Pro, jax-mps 0.10.9, safe dispatch):
T= 25 no-stack 1.6 ms stacked 139.1 ms stack cost/step 5.50 ms
T= 50 no-stack 2.0 ms stacked 507.5 ms stack cost/step 10.11 ms
T= 100 no-stack 3.4 ms stacked 1986.7 ms stack cost/step 19.83 ms
T= 200 no-stack 6.4 ms stacked 7809.1 ms stack cost/step 39.01 ms
Per-step cost doubles every time T doubles, so the total is ~T². The stacked output here is 100 × 1e6 × 4 B = 0.4 GB at T=100; ~2 s to produce it implies roughly 20 GB of traffic. On CPU the same MWE is linear.
Working hypothesis: the scan's output accumulator is a loop-carried value of the while body, and each iteration's stablehlo.dynamic_update_slice materializes a fresh copy of the whole (T, N) buffer instead of updating in place — MLX arrays are immutable and the primitive holds a live reference to the previous iteration's accumulator, so MLX's donation/in-place path can't engage. If that's right, the fix is buffer aliasing for loop-carried accumulators inside WhileLoopPrimitive, which is executor work rather than a lowering change. For MCMC specifically this compounds #196: chains of realistic length are currently dominated by this rather than by RNG.
Happy to help test a fix at large N/T; I have an M3 Pro doing particle-filtering workloads where this is the binding constraint.
jax: 0.10.2
jaxlib: 0.10.2
jax-mps: 0.10.9
Python: 3.13.9
macOS: 26.2 (25C56)
Chip: Apple M3 Pro
Memory: 36 GB
Backend: mps
Follow-up in the same area as #134 and #193/#196. The
WhileLoopPrimitivefrom #134 fixed the graph-growth crash, but a scan that stacks outputs — which is how JAX collects MCMC samples, RNN hidden states, filter histories — now pays a per-iteration cost that grows with T, so total time is quadratic. A scan that carries the same state without stacking is flat.MWE:
Results (M3 Pro, jax-mps 0.10.9, safe dispatch):
Per-step cost doubles every time T doubles, so the total is ~T². The stacked output here is 100 × 1e6 × 4 B = 0.4 GB at T=100; ~2 s to produce it implies roughly 20 GB of traffic. On CPU the same MWE is linear.
Working hypothesis: the scan's output accumulator is a loop-carried value of the while body, and each iteration's
stablehlo.dynamic_update_slicematerializes a fresh copy of the whole (T, N) buffer instead of updating in place — MLX arrays are immutable and the primitive holds a live reference to the previous iteration's accumulator, so MLX's donation/in-place path can't engage. If that's right, the fix is buffer aliasing for loop-carried accumulators insideWhileLoopPrimitive, which is executor work rather than a lowering change. For MCMC specifically this compounds #196: chains of realistic length are currently dominated by this rather than by RNG.Happy to help test a fix at large N/T; I have an M3 Pro doing particle-filtering workloads where this is the binding constraint.