Optax: optimizers

Optax provides gradient processing and optimization functionality on top of JAX, including optimizers and losses.

Main Optax site: optax.readthedocs.io

If you’d like to learn more about Optax beyond what’s covered in the Getting started with JAX for ML guide, we recommend starting with the Optax getting started guide.