Autoregressive rollout in jax February 26, 2025