Glossary term
Glossary term
Infrastructure and Serving
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.
Created for this library
A research team uses Optax for composable optimizer recipes in JAX, mixing learning rate schedules and gradient clipping cleanly.
An ML platform team adopts Optax to standardize optimizer configurations across JAX-based training jobs.
An ML researcher uses Optax to combine custom clipping and decay schedules for a new training recipe.
Definition source: Google for Developers Machine Learning Glossary | Creative Commons Attribution 4.0 License