# Before (Python for loop - gets unrolled)
for t in range(n):
    key, subkey = jax.random.split(key)
    Z = jax.random.normal(subkey, (2, M))
    s = s + μ + jnp.exp(h) * Z[0, :]
    h = ρ * h + ν * Z[1, :]

# After (JAX fori_loop - compiled efficiently)
def update(i, state):
    s, h, key = state
    key, subkey = jax.random.split(key)
    Z = jax.random.normal(subkey, (2, M))
    s = s + μ + jnp.exp(h) * Z[0, :]
    h = ρ * h + ν * Z[1, :]
    return s, h, key

s, h, key = jax.lax.fori_loop(0, n, update, (s, h, key))