optax

Why we included this project

Training models in JAX means the optimizer often matters as much as the network itself, and Optax is where most practitioners land. It ships the standard gradient-processing pieces (Adam, SGD with momentum, AdaGrad, and others) as small composable building blocks instead of a rigid training loop. You can plug in a learning-rate schedule or combine gradient transformations into a custom optimizer without rewriting your training code. The implementations are well tested and written to match the standard equations, which makes the math easier to trust and adapt for research experiments or production pipelines. For anyone already in the JAX ecosystem, it is a pragmatic foundation for getting gradient updates right.