jax.checkpoint ============== .. currentmodule:: jax .. autofunction:: checkpoint