Optax is a gradient processing and optimization library for JAX.

Optax is designed to facilitate research by providing building blocks that can be easily recombined in custom ways.

The goals are to: - Provide simple, well-tested, efficient implementations of core components. - Improve research productivity by enabling to easily combine low level ingredients into custom optimisers (or other gradient processing components). - Accelerate adoption of new ideas by making it easy for anyone to contribute.