Nonlinear solvers that actually compose

What it does

Optimistix is a JAX library for nonlinear numerical methods—root finding, minimisation, fixed-point iteration, and least-squares solving. It tackles problems where simple gradient descent isn’t enough, and it speaks native JAX: autodiff, PyTrees, GPU/TPU dispatch, and all. If you are already using Equinox, it slots in without friction.

The interesting bit

The library treats solver components like interchangeable parts. You can pair a BFGS quadratic approximation with a dogleg descent path and a trust-region updater, or automatically reformulate a root-finding problem as a least-squares task and feed it to a minimiser. That modularity is unusual in numerical computing, where optimisers are typically all-or-nothing black boxes.

Key highlights

• Problem-type interoperability: autoconvert root-finding to least-squares to minimisation without manual rewriting.
• Composable internals: swap quadratic models, descent paths, and trust-region strategies independently.
• Native PyTree state: solver state matches JAX’s nested tree structures, so you are not stuck flattening parameters.
• Plays nicely with Optax and the broader JAX scientific stack (Equinox, Lineax, Diffrax).

Verdict

Worth a look if you are doing implicit layers, inverse problems, or any JAX work that needs Newton-style methods rather than just gradient descent. Skip it if your workload is purely first-order optimisation—Optax already has you covered.