Optax: Composable Gradient Transformations for JAX
Optax is a gradient processing and optimization library for JAX.
At a glance
- What is it?
- Optax is DeepMind's library of gradient transformations and optimizers for JAX, built around small combinable pieces rather than monolithic optimizer classes. It fits teams already committed to JAX who need custom update rules; it is not a framework-agnostic optimizer collection.
- Who is it for?
- Adopt Optax if your training loop is already written in JAX and you need to assemble custom update rules from small parts, or if you want standard optimizers such as optax.adam with a consistent init/update interface. Do not adopt it if you are working in PyTorch or TensorFlow, or if you want a single turnkey training framework rather than building blocks.
- Can I use it commercially?
- Yes. Apache-2.0 is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 2 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The problem Optax solves: gradient processing as separate building blocks
Most deep learning frameworks ship optimizers as self-contained classes. Adam is one object, SGD is another, and combining them (say, clipping gradients before Adam, or applying weight decay only to certain parameter groups) usually means either subclassing or waiting for the framework to add the combination you want. Optax takes a different position. The README states that the library is designed to facilitate research by providing building blocks that can be easily recombined in custom ways, and that the project favors focusing on small composable building blocks that can be effectively combined into custom solutions. That single design decision explains most of what follows, including both the flexibility and the extra code you write.
The audience is narrower than a general optimizer library. Optax is a JAX library. Its examples use jnp arrays and jax.grad, and its documentation lives alongside the JAX ecosystem. If your model is not in JAX, Optax is not the tool. Within JAX, the target user is someone doing research or non-standard training who needs an update rule that does not exist as a finished class, or who wants to read an optimizer's implementation and match it against the equations in a paper. The README is explicit that implementations prioritize readability and structuring code to match standard equations over code reuse. That is a deliberate trade: less shared machinery, more code that reads like the paper it implements.
How Optax structures an update: init, update, apply_updates
The mechanism is a three-step contract, and it is visible directly in the README quickstart. First, you construct an optimizer by calling a function such as optax.adam(learning_rate). Second, you initialize optimizer state against your parameters with optimizer.init(params). The README notes that the resulting opt_state contains statistics for the optimizer, which for Adam means the moment accumulators. Third, inside the training loop you call optimizer.update(grads, opt_state), which returns a tuple of updates and a new opt_state, and then you apply those updates with optax.apply_updates(params, updates).
What makes this different from a typical optimizer class is that the optimizer is a transformation, not a mutable object holding parameters. State is passed in and returned explicitly, which is what makes the pattern compatible with JAX's functional style and with jit compilation. The README's example declares params as a dictionary, {'w': jnp.ones((num_weights,))}, and the loss function indexes into it: optax.l2_loss(params['w'].dot(x), y). So the parameter pytree is part of the interface, and the optimizer state is built to match its structure.
Because the pieces are transformations, the library also ships loss functions alongside optimizers. The quickstart uses optax.l2_loss for mean squared error, and the documentation links to separate API listings for optimizers and for losses. The gradient transformation framing is the important part: clipping, scaling and the optimizer step are all things that can be expressed in the same vocabulary, which is why they can be chained.
Getting Optax installed and running the quickstart loop
Installation is a single command for the released version, per the README:
pip install optax
For the development version from GitHub:
pip install git+https://github.com/google-deepmind/optax.git
The quickstart then follows the three-step contract. You create the optimizer and state:
optimizer = optax.adam(learning_rate) params = {'w': jnp.ones((num_weights,))} opt_state = optimizer.init(params)
You compute gradients with jax.grad over a loss built from an Optax loss function:
compute_loss = lambda params, x, y: optax.l2_loss(params['w'].dot(x), y) grads = jax.grad(compute_loss)(params, xs, ys)
And you close the loop:
updates, opt_state = optimizer.update(grads, opt_state) params = optax.apply_updates(params, updates)
The README points to a getting started notebook for continuing past this snippet. For contributors, the repository layout adds a few more commands. The source is cloned with git clone https://github.com/google-deepmind/optax.git, tests run via sh test.sh, and documentation builds after installing the docs extra with pip install -e ".[docs]" followed by make html -C docs. The README also asks that anyone adding a feature such as a new optimizer open an issue first, which is a real constraint on contribution workflow rather than a formality.
Where Optax puts work back on you
The composability that makes Optax flexible also means the training loop is yours. Nothing in the material shows a Trainer, a fit method, or a checkpointing utility. You call update and apply_updates yourself, and you keep opt_state alive across steps and across checkpoint boundaries. If your team is used to frameworks that serialize optimizer state for you, that is a piece of plumbing you now own.
The second cost is API surface. The README links an optimizers API page and a losses API page rather than enumerating what exists, so the accurate statement is that the set of available optimizers is defined by that listing, not by the README text. Anyone evaluating Optax for a specific algorithm should check the current listing rather than assume coverage. The README does say the library contains implementations of many popular optimizers, which is a claim about breadth, not a guarantee about the one you need.
The third cost is JAX itself. Optax inherits JAX's constraints: functional purity, explicit state, and the requirement that the loss be differentiable by jax.grad. The README's own example is a linear model, which sidesteps the harder case of stateful layers in the parameter pytree. If your model code is not already written in this style, adopting Optax means adopting that style first. That is not a defect in Optax, but it is the reason a PyTorch user will find the quickstart unfamiliar rather than merely different.
The nearest alternative in the JAX ecosystem
The README itself lists optimization-adjacent libraries in JAX, and the first is optimistix, described as providing nonlinear solvers: root finding, minimisation, fixed points, and least squares. The difference in approach is worth stating precisely. Optax operates on gradients. You bring the gradients, produced by jax.grad, and Optax transforms them into parameter updates. Optimistix, by the README's description, solves the optimisation problem itself: you give it a function and it handles root finding, minimisation, fixed points or least squares. That places it in a different slot in the stack. If you have a loss and want a minimiser without writing a gradient loop, the optimistix description matches that need. If you have gradients from a neural network training step and want to control how they become updates, Optax matches that need.
The README also lists matfree, described as matrix free methods useful to study curvature dynamics in deep learning. That is a research tool rather than a drop-in optimizer replacement, and the README presents it as such. The honest summary is that Optax is not competing with a general-purpose optimizer library in another framework; it is one layer in the JAX stack, and the neighbouring layers do different jobs.
Maintenance, releases and licence
The release cadence visible in the supplied material is roughly two to three releases per year: v0.2.6 in September 2025, v0.2.7 in February 2026, and v0.2.8 in March 2026, with the last push to the repository in September 2026. The version numbers are all 0.2.x, which signals that the project does not present itself as API-stable at the 1.0 level. For adopters that means pinning a version and reading release notes before upgrading is a reasonable posture, though the material here does not detail what changed between those releases, so the specific breakage risk cannot be assessed from this page alone. The repository is not archived and the README states that issues and pull requests are welcome, with the caveat that new features such as optimizers should be proposed in an issue first.
On licensing, the repository is Apache-2.0. That is a permissive licence, and it is the same licence family used across much of the surrounding ecosystem, but the material here does not include the full licence text or any statement about patent grants or attribution requirements. Anyone distributing a product that embeds Optax should read the LICENSE file in the repository and, if the stakes are high, get their own legal review. Nothing in this article should be read as legal advice. The citation block in the README asks users to cite the DeepMind JAX Ecosystem paper, which is an academic convention rather than a licence term.
Who should adopt Optax, and what to check first
Optax fits a specific situation: your training code is already in JAX, you understand that the optimizer is a transformation with explicit state, and you either need a standard optimizer such as optax.adam or you need to build a custom update rule from smaller pieces. The README's stated goal of enabling researchers to combine low-level ingredients into custom optimizers is the clearest statement of intent, and the three-function contract of init, update and apply_updates is small enough to learn in an afternoon.
It fits poorly in three cases. If your model is not in JAX, the quickstart will not translate. If you want a training framework that owns the loop, checkpoints and logging, Optax deliberately does not provide that layer. And if your need is solving an optimisation problem rather than transforming gradients, the README's own pointer to optimistix is the better starting point.
Before adopting, verify the current optimizer listing against the algorithm you need, since the README does not enumerate it. Verify that you are prepared to own optimizer state persistence across checkpoints, because the material shows no utility for it. And verify the Apache-2.0 terms against your distribution plans by reading the LICENSE file directly. The library is small, composable and readable by design; the cost of that design is that the surrounding loop is yours to write and maintain.
Editorial conclusion
Adopt Optax if your training loop is already written in JAX and you need to assemble custom update rules from small parts, or if you want standard optimizers such as optax.adam with a consistent init/update interface. Do not adopt it if you are working in PyTorch or TensorFlow, or if you want a single turnkey training framework rather than building blocks. Before committing, verify three things: that the optimizer you need exists in the current API listing, that you are comfortable maintaining the update loop yourself, and that the Apache-2.0 licence terms fit your distribution model.
Community notes