Optax is a gradient processing and optimization library for JAX. It provides composable building blocks for implementing optimizers and training routines for parametric models such as deep neural networks.
A gradient processing and optimization library for JAX. Optax facilitates research by providing building blocks that can be recombined in custom ways to optimize parametric models such as deep neural networks. Other goals include: - Providing readable, well-tested, efficient implementations of core components. - Improving productivity by making it possible to combine low level ingredients into custom optimizers (or other gradient processing components). - Accelerating adoption of new ideas by making it easy for anyone to contribute.