Optimistix: nonlinear solvers that compose inside JAX

Four solver families, one library entry point

Optimistix covers root finding, minimisation, fixed points and least squares. Those are usually four separate packages with four separate APIs, and the claim here is that they are one library that treats them as convertible views of each other.

The quick example shows what that looks like in practice. The problem posed is the implicit Euler step for the ODE dy/dt = tanh(y(t)), which requires finding a value y1 satisfying y1 = y0 + tanh(y1) * dt. That is a fixed point problem, and it is solved with `optx.fixed_point(fn, solver, y0)` where `fn` returns `y0 + jnp.tanh(y) * dt` and the solver is a Newton instance with both a relative and an absolute tolerance set to 1e-5. The result comes back as `sol.value`, and the assertion in the comment is that it equals `fn(y1)`.

The signature matters more than the example. A function taking the state and an argument tuple, a solver object, and an initial value, returning a solution with a value attribute, is the same shape across the four families. Learning one is learning all of them.

Root finding autoconverts into least squares

The interoperability claim is specific: a root finding problem can be autoconverted into a least squares problem and then solved with a minimisation algorithm.

That is a bigger deal than it sounds. Root finding normally demands derivatives of the residual and a convergence criterion tuned to the residual, while least squares can use gradient information and reuse the optimisers already written for fitting. Converting between them means you can pick the solver family for numerical reasons rather than for historical reasons, and it means the least squares machinery is exercised by root finding workloads it was not explicitly written for.

The same modularity applies to the optimisers themselves, and the README gives the composition as an example: a BFGS quadratic bowl, a dogleg descent path, and a trust region update. Those are three separate choices. A quadratic approximation, a way to move inside it, and a mechanism for resizing the region of trust are the parts that solvers usually fuse together, and treating them as swappable is what lets you pair a dogleg path with something other than BFGS.

Neither mechanism is a claim that the default is wrong. They are claims that the default is not the only option.

State is a PyTree, so the JAX compiler carries it

The solver state is a PyTree, which is the decision that makes everything else in the list free.

Because the state is a tree rather than a tensor, a solver can carry whatever its algorithm needs: a current iterate, a history, a Jacobian, timing counters. JAX already knows how to map, differentiate, trace and place PyTrees, so none of that has to be reimplemented for the solver, and the library inherits the rest of the JAX feature list without mentioning it: automatic differentiation, automatic parallelisation across devices, and GPU and TPU support.

The listed benefits are fast compilation and runtimes, and interoperability with Optax. The second one is the practical bridge: if your objective needs a first-order optimiser to get anywhere near a solution, Optimistix will use one from Optax rather than asking you to write a descent loop.

Installation is a single line and the floor is Python 3.11:

pip install optimistix

JAX 0.7.0 and 0.7.1 are explicitly rejected

The dependency declaration carries one detail that will save you an afternoon. The JAX requirement is `jax>=0.4.38` with two exclusions written directly into it, `!=0.7.0` and `!=0.7.1`.

Version pins like that are written for a reason, and neither of these two releases is a version you can simply upgrade into. If your environment resolves to one of them, installation fails at resolution time rather than producing a solver that misbehaves later, which is the better of the two failure modes.

The rest of the runtime dependencies are four: jaxtyping for shape and dtype annotations on arrays, lineax for linear solvers, equinox for the modelling layer, and typing_extensions. Lineax being a direct dependency rather than an optional extra is the interesting one: a nonlinear method that converges to a root often wants a linear solve at each step, and having it available without a separate install removes a class of version mismatch.

The keyword list includes levenberg-marquardt, which tells you where the effort went in the least squares side of the library.

Alpha status with three pinned tool groups

The project classifies itself as Development Status 3, Alpha, and targets an unusually broad audience for a scientific library: developers, science and research, information technology, and financial and insurance industry. The last of those is not decorative, since root finding and least squares are the numerical core of a calibration pipeline.

Development dependencies are split into three groups and every one is pinned to an exact version rather than a range. The dev group holds prek, pyright, ruff and toml-sort. The docs group is larger and includes mkdocs with the material theme, mkdocstrings with the Python handler, griffe and hippogriffe for introspection, mkdocs-ipynb so notebooks are rendered, pygments, and pymdown extensions.

The test group adds beartype for runtime shape checking, pytest, diffrax, optax, jaxlib, sif2jax, fire and matplotlib.

That combination is worth reading as a statement about how the library verifies itself. Diffrax in the test group means the solvers are checked against an independent differential equation solver rather than against themselves, and beartype means shape errors surface at the boundary instead of inside a compiled step.

The release history is unhurried: v0.0.10 in December 2024, v0.0.11 in October 2025, and v0.1.0 in February 2026, with the last push to main dated 2026-08-11.

The see-also list is a map of what Optimistix does not do

The companion projects listed in the README are more informative than a feature list, because they partition the numerical territory.

Optax is described as first-order gradient optimisers, that is SGD and Adam and their relatives. Optimistix is the other half: nonlinear solvers. Lineax is linear solvers, and Diffrax is numerical differential equation solvers. Orbax covers checkpointing across asynchronous, multi-host and multi-device runs. paramax is parameterisations and constraints for PyTrees, and jaxtyping is the type annotation layer.

Two entries are flagged as neighbours rather than dependencies. sympy2jax converts SymPy expressions to JAX so symbolic expressions can be trained by gradient descent, and PySR is symbolic regression, marked as a non-JAX honourable mention.

The adjacency matters for a specific reason. An implicit ODE step like the one in the quick example is a fixed point problem, and Diffrax, which solves ODEs, has to solve fixed points internally too. Having both libraries available means you can use Diffrax where the integration is the hard part and Optimistix where the root find is the hard part, without either one reimplementing the other's strength.

The documentation lives at docs.kidger.site/optimistix, and the paper is on arXiv, with a BibTeX entry in the README for citing it.

Editorial conclusion

Use Optimistix when your problem is nonlinear and you want it inside a JAX pipeline: implicit ODE steps, fixed points, least squares against an autodiff model, or a root find you would rather express as a minimisation. Do not reach for it to replace a first-order optimiser, since that is Optax's job and Optimistix only interoperates with it. Before you pin a JAX version, read the exclusion list, because 0.7.0 and 0.7.1 are explicitly rejected, and treat the alpha classification as a signal about the API surface rather than about the maths.

Frequently asked questions

Official sources

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