Practical overview
When it fits
Use Optax when a JAX training or research workflow needs reusable optimizer and gradient-processing components, rather than a hosted service that manages the full training process.
Who it is for
- JAX researchers composing or adapting optimization methods for experiments.
- Machine-learning developers who need optimizer state and gradient-update components in Python training code.
Documented uses
- Choose an implemented optimizer such as Adam for a JAX model.
- Turn gradients from a JAX-differentiable loss into parameter updates.
- Combine low-level optimization components into a custom optimizer or gradient-processing workflow.
Requirements & limits
The repository is public, identifies the project as open source, and lists an Apache-2.0 license. The README documents installation from PyPI or directly from GitHub.
- Optax provides optimization components; the surrounding model, data and JAX loss workflow remain the user's responsibility.
- Its design favors small building blocks that can be recombined, rather than requiring one fixed optimizer abstraction.
- The documented quickstart relies on JAX differentiation and uses optimizer initialization, update and parameter-application stages.
Check current requirements in the official documentation. Source-based guidance does not mean OpenViST has run or independently tested this project.
Install Optax and connect an optimizer to JAX gradients
Add an Optax optimizer to a JAX update workflow using the documented quickstart pattern.
Before you start
- Python environment with JAX code
- Model parameters and a loss function differentiable by JAX
- A learning-rate value for the optimizer
Steps
Install the released package with the command below.
pip install optaxCreate an optimizer, such as
optax.adam(learning_rate), and initialize its state withoptimizer.init(params).Use
jax.gradon a JAX-differentiable loss to obtain gradients.Call
optimizer.update(grads, opt_state)and apply the resulting updates withoptax.apply_updates(params, updates).
Expected outcome
A JAX code workflow that converts loss gradients into parameter updates while carrying optimizer state.
Guide limitations
- The example assumes the user supplies the model parameters, data and differentiable loss.
- The documented steps describe the optimizer update pattern, not a complete model-training or deployment system.
Sources & review
Source record details
- Source availability
- Active
- Sources checked
- 13 October 2026
- Original website title
- google-deepmind/optax · Official homepage title · retained in its source language
Related options
Optimistix — A JAX optimization-adjacent library described by the Optax README for nonlinear solvers, including root finding, minimisation, fixed points and least squares.
Matfree — A JAX optimization-adjacent library described by the Optax README for matrix-free methods useful in studying curvature dynamics in deep learning.