JAXopt: Differentiable Optimizers in JAX, Now Archived by Its Maintainers

The gap JAXopt was built to fill between JAX and an optimizer library

Plain JAX gives you autodiff, vmap and jit, but it does not give you a solver. If you want L-BFGS, a proximal gradient step, or the solution of a fixed-point equation, you write the loop yourself. JAXopt's stated purpose is to supply those loops as JAX-native objects: hardware accelerated across CPU, GPU and TPU, batchable through vmap, and differentiable with respect to the inputs of the problem rather than only the parameters of a model. That last property is the one that separates it from a training-loop optimizer. The intended user is someone doing differentiable programming or bi-level optimization, where an inner problem is solved and the outer loss depends on that solution. The topics list on the repository says the same thing: bi-level, differentiable-programming, deep-learning, jax, optimization.

Implicit differentiation is the part that is hard to replace

The README describes two ways to differentiate an optimization solution: implicitly, or via autodiff of unrolled algorithm iterations. Unrolling is the approach you get for free from any solver written in JAX, and it costs memory proportional to the number of iterations because every intermediate state stays on the tape. The implicit route instead differentiates the optimality condition, so the backward pass does not replay the solver. The repository cites the paper Efficient and Modular Implicit Differentiation (arXiv:2105.15183) as the description of that framework. This is a real architectural difference, not a marketing one: it changes the memory profile of the backward pass and it changes what happens when the solver has not converged, because an implicit gradient assumes the solution satisfies the fixed-point condition. If your inner solver terminates early, the gradient you get is the gradient of the condition you did not actually satisfy.

Batching and hardware acceleration come from JAX, not from JAXopt

The three adjectives in the project description are worth separating. Hardware acceleration and vmap batching are inherited properties: because the solvers are written in JAX primitives, vmap can map over a batch of independent problems and the same code runs on GPU or TPU without a separate implementation. JAXopt does not add a scheduler or a device abstraction on top. The practical consequence is that batching is only as good as the solver's control flow. A solver whose iteration count depends on data will still work under vmap, but the mapped computation runs until the slowest instance in the batch finishes, so a batch of problems with wildly different conditioning is inefficient in a way that is inherent to the approach rather than a bug in JAXopt. The README does not discuss this cost, and that is a gap worth knowing about before you build a pipeline around batched solves.

Installation and the version you should actually pin

The README gives three install paths. The released version is `pip install jaxopt`. The development version is `pip install git+https://github.com/google/jaxopt`. Installing from a source checkout is `python setup.py install`. For a project that the README declares unmaintained, the first command is the only one that makes sense for anything reproducible, and even then you should pin the release identifier rather than take whatever pip resolves today. The releases listed are jaxopt-v0.8.3 (2024-01-10), jaxopt-v0.8.4 (2025-04-10) and jaxopt-v0.8.5 (2025-04-14). The gap between 0.8.3 and 0.8.4 is roughly fifteen months, which tells you the cadence had already slowed well before the archive notice. Pinning `jaxopt==0.8.5` in a requirements file is the difference between a dependency and a moving target, because the git install path will track `main` and `main` is not being developed.

The maintenance status is the headline, not a footnote

The README's Status section says plainly that JAXopt is no longer maintained nor developed, and directs readers to the JAX website for alternatives. It also states that some features, specifically losses, projections and the lbfgs optimizer, have been ported into optax. That sentence is the most actionable line in the document: it is a partial migration map. Anything in that list has a maintained home. Anything outside it does not, as far as the supplied material shows. The disclaimer adds that JAXopt was an open source project maintained by a team in Google Research and is not an official Google product, which matters if you were assuming a support commitment came with the name. There is no deprecation timeline, no stated end-of-life date, and no compatibility guarantee for future JAX releases. A pinned install will keep working until a JAX change breaks it, and when that happens the fix is yours.

Where optax is the better answer, and where it is not

Optax is the maintained alternative named in the README, and the difference in approach is not cosmetic. Optax is built around gradient transformations: you compose a chain of updates that are applied to parameters during a training loop. JAXopt is built around solving a problem to a stopping criterion and then differentiating the solution. If what you want is Adam, SGD with a schedule, or L-BFGS applied to model parameters during training, optax covers it and the README says the lbfgs optimizer has already moved there. If what you want is to solve an inner optimization or a fixed-point problem and backpropagate through the result, the transformation model does not express that, and that is the slice JAXopt still occupies. The honest framing is that optax replaced part of JAXopt and left the implicit differentiation machinery behind.

Licence, upgrades and what a frozen dependency costs

The licence is Apache-2.0, which permits commercial use, modification and redistribution provided you keep the notices and state changes, and it includes a patent grant. That is a permissive licence and it does not obligate you to publish your own code. This is a description of the licence text, not legal advice; if the patent grant or notice requirements matter to your organisation, have counsel read it. The upgrade cost is the real expense. Because the project is unmaintained, there is no upstream to fix a breakage when a new JAX release changes a primitive, and no release notes to tell you what changed. Your options are to vendor the specific solver files you use, to fork them under the same Apache-2.0 terms, or to rewrite the call site against whatever the JAX website now recommends. Vendoring a handful of solver modules is cheaper than forking the repository, and it makes the frozen surface explicit in your own tree.

Who should still use it, and what to check before you do

Use JAXopt if you are doing differentiable programming where the gradient flows through a solved problem, you specifically need the implicit differentiation path described in arXiv:2105.15183, and you are willing to pin `jaxopt==0.8.5` and treat the solver code as code you own. Do not use it as the optimizer for a standard training loop, because optax is maintained and the README says the lbfgs optimizer already lives there. Do not use it if you need a support channel or a compatibility promise, because the README offers neither. Before writing any code, check whether your solver is one of the three features the README says were ported to optax, and check whether the implicit differentiation framework has an equivalent on the JAX website the README points to. If neither check resolves in your favour, the decision is between vendoring the solver you need and writing the optimality-condition gradient yourself, and the second option is often less work than it sounds.

Editorial conclusion

Adopt JAXopt only if you need its implicit differentiation framework, meaning the custom_root and fixed_point solvers that let you differentiate through a solution without unrolling iterations, and you can pin the version. Do not adopt it as a general-purpose optimizer library for new work, because the README states it is no longer maintained or developed and points to the JAX website for alternatives. Before committing, verify whether the solver you need has a direct equivalent in optax, since the README says losses, projections and the lbfgs optimizer have already been ported there. If your solver is not in that list, assume you will be maintaining the call site yourself.

Official sources

1. google/jaxopt on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases