Optimistix: nonlinear solvers that live inside JAX transformations

The gap Optimistix fills between Optax and SciPy

JAX ships autodiff, autoparallelism and GPU/TPU support, but it does not ship nonlinear solvers. Optax covers first-order gradient methods such as SGD and Adam, which is the right tool when you are descending a loss surface with many parameters. It is the wrong tool when the thing you want is a root: a value y such that y equals some function of y. The README's own example is exactly that case. Solving dy/dt = tanh(y(t)) with implicit Euler means finding y1 such that y1 = y0 + tanh(y1) * dt, and no amount of gradient descent answers that question directly. Optimistix targets that class of problem: root finding, minimisation, fixed points and least squares, written so the solve can happen inside JAX transformations. The audience is scientific computing and probabilistic modelling work in JAX, not general deep learning. The repository topics list deep-learning and neural-networks, but the README's own framing and its See also list place it next to Diffrax, Lineax and BlackJAX rather than next to a training loop.

Modular solvers and automatic problem conversion

Two design decisions separate Optimistix from a typical solver package. The first is problem conversion. The README states that solvers are interoperable, giving the example of autoconverting a root find problem into a least squares problem and then solving it with a minimisation algorithm. That matters because root finding and least squares are usually separate entry points in numerical libraries, and the user is left to reformulate by hand. The second is modularity at the level of the iteration itself. The README describes using a BFGS quadratic bowl with a dogleg descent path and a trust region update, which means the quadratic model, the step acceptance rule and the trust region radius update are separate objects you can combine. A conventional library hands you BFGS or dogleg as a finished method. The state is a PyTree, so the solver's internal state is a JAX-native structure rather than a flat array or an opaque C struct, which is what allows the whole solve to be traced. Interoperability with Optax is listed as a feature, suggesting the two are meant to be used together rather than as substitutes.

Installing and running the fixed-point example

Installation is a single command, pip install optimistix, and the README states the requirement as Python 3.11 or newer. That floor is worth noting: it is higher than the 3.9 or 3.10 baseline many scientific stacks still pin. The quick example is short enough to reproduce in full. It imports jax.numpy as jnp and optimistix as optx, defines y0 = jnp.array(1.) and dt = jnp.array(0.1), then defines fn(y, args) returning y0 + jnp.tanh(y) * dt. The solver is constructed as optx.Newton(rtol=1e-5, atol=1e-5), and the call is sol = optx.fixed_point(fn, solver, y0). The result is read from sol.value, which satisfies y1 == fn(y1). The argument order matters: the function comes first, then the solver, then the initial guess. The function signature takes (y, args) even though args is unused here, which is the convention you will carry into every other problem type. Tolerance is set on the solver object, not passed to the solve call, so a solver instance is configured once and reused.

Where the modular design costs you

The same flexibility that lets you combine a quadratic bowl with a dogleg path and a trust region update also means the library cannot guarantee that every combination is sensible. The README presents the mixing as a feature and does not publish a compatibility matrix in the material available here. If you assemble a method from parts, the convergence behaviour is your responsibility, not the library's. The second limitation is structural. Because the state is a PyTree and the solver runs under JAX tracing, anything in your residual function that is not traceable will fail, and control flow that depends on concrete values will either be traced incorrectly or force a recompile. This is the standard JAX tax, but it lands harder on solvers than on training loops, because solver iterations are inherently sequential and data-dependent. If your problem is a single small root find called once at the top of a script, Optimistix is the wrong tool. SciPy's root and least_squares functions will be faster to write and will not ask you to think about tracing, and the JAX overhead buys you nothing when there is no batch and no jit.

Optimistix versus Optax, and versus SciPy

The comparison with Optax is the one the README invites, and the difference is not a matter of quality. Optax solves one problem: descending a scalar loss with first-order gradient information, with SGD, Adam and similar rules. Optimistix solves a different one: finding points that satisfy equations, including cases where no loss function exists to descend. The README lists interoperability with Optax as a feature, which is the honest description of the relationship. A least squares fit in Optimistix and an Adam run in Optax are not competing answers to the same question. The comparison with SciPy is sharper. SciPy's solvers are mature, widely exercised and callable from ordinary Python, with no tracing constraints and no PyTree requirement. Optimistix trades that maturity for composability inside JAX: autodiff, GPU and TPU execution, vmap over a batch of independent solves, and the ability to differentiate through the solve. If you need any of those, SciPy cannot follow you. If you need none of them, SciPy is the shorter path.

Version history and what an upgrade costs

The release record is thin and unevenly spaced. v0.0.10 landed in December 2024, v0.0.11 in October 2025, and v0.1.0 in February 2026, with the last push to the default branch in August 2026. That is roughly a year between the first two of those releases and about four months to the 0.1.0 tag. The jump from 0.0.11 to 0.1.0 is the one to watch: in a project that has sat in the 0.0.x series, a 0.1.0 tag is the conventional signal that the maintainers consider something settled, though the README does not state what changed. Because Optimistix is a JAX library, its real upgrade cost is not the Optimistix version but the JAX version underneath it. Equinox supplies the PyTree machinery, and JAX changes to tracing, transformation internals or dtype promotion can force retracing or break a solver step even when the Optimistix version is unchanged. Pin JAX and Equinox together with Optimistix, and treat a JAX bump as a change that needs its own test run rather than a routine dependency refresh.

Licence and maintenance posture

Optimistix is Apache-2.0, which permits commercial use, modification and redistribution provided the licence and notices are preserved, and it includes a patent grant. That is a permissive choice, and it is the same licence family used across much of the JAX ecosystem, so it is unlikely to create friction in a corporate codebase. This is a description of the licence text, not legal advice; if your organisation has specific requirements around patent clauses or notice files, have counsel read it. On maintenance, the README names the primary author as Jason Rader and states that the project is co-maintained by Johanna Haffner. Two named maintainers is a real bus factor, and it is the number to weigh when you decide how deeply to depend on the library. The citation entry points to arXiv:2402.09983, so there is a paper describing the design if you need to evaluate the method choices rather than the API.

Editorial conclusion

Adopt Optimistix if your solver step has to sit inside a jit, a vmap, or a gradient computation, and if you want to assemble a method from parts rather than pick a monolith. Do not adopt it for plain neural network training, where Optax is the documented first-order path, and do not expect it to replace a mature SciPy workflow without rewriting your problem as a PyTree. Verify two things before committing: that the solver combination you want is actually exposed for your problem class, and that your pinned JAX version still compiles the Equinox PyTree state without retracing.

Official sources

1. Issues
2. License: Apache-2.0
3. patrick-kidger/optimistix on GitHub
4. README
5. Releases