def do_multistep(state, num_steps):
    """Perform multiple model steps back-to-back."""
    return jax.lax.fori_loop(
        0, num_steps, lambda _, s: shallow_water_step(s, False), state
    )
Perform multiple model steps back-to-back.
examples/shallow_water.py