A gradient processing and optimisation library in JAX.
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.
