Optax — библиотека обработки градиентов и оптимизации для JAX. Она предоставляет компонуемые строительные блоки для реализации оптимизаторов и процедур обучения параметрических моделей, таких как глубокие нейронные сети.
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.