optimistix

What it is and what it does

Optimistix is a JAX library for solving nonlinear problems: finding roots, minimizing functions, computing fixed points, and solving least-squares problems. It wraps these problem types into a unified interface so you can, for example, convert a root-finding problem to a least-squares problem and solve it with a minimization algorithm. The library is built around modular solvers—you can compose descent paths (like dogleg), update rules (like trust region), and quadratic models (like BFGS) to construct custom optimizers.

Because it's built on JAX, Optimistix inherits autodifferentiation, automatic parallelization, and GPU/TPU support. It works with PyTree-based state, integrates with the Equinox neural-network library, and can interoperate with Optax optimizers. The library is in alpha (Development Status 3) and actively maintained, with documentation available online.

Use it for:

• Solve implicit differential equations (e.g., implicit Euler) by finding fixed points of the update map.
• Minimize loss functions in neural networks or scientific models using custom composed solvers.
• Find roots of nonlinear systems for physics simulations or engineering problems.
• Solve least-squares problems for parameter fitting or inverse problems.
• Combine multiple solver types in a single JAX-compiled pipeline with autodiff.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Optimistix provides nonlinear solvers for root finding, minimization, fixed-point problems, and least-squares fitting, built on JAX with support for autodiff, GPU/TPU acceleration, and PyTree-based state.

Yes, if you work with JAX and need nonlinear solvers. The library is actively maintained, has no known vulnerabilities, installs with low friction, and offers a composable solver interface tailored to JAX's autodiff and GPU/TPU ecosystem. The alpha status means the API may change, but the maintenance signal is strong. Not necessary if you only need first-order gradient optimizers (use Optax instead) or if you are not already in the JAX ecosystem.

Install

optimistix on PyPI

pip install optimistix

uv add optimistix

poetry add optimistix

Installing optimistix

Before you install

Low friction: pure Python wheel with five runtime dependencies (equinox, jax, jaxtyping, lineax, typing-extensions). Actively maintained with recent commits and no known vulnerabilities. Requires Python 3.11+.

License in practice

Apache 2.0 permissive license allows commercial and derivative use with minimal restrictions—include a copy of the license and note any modifications to the source.

Quickstart

pip install optimistix

import jax.numpy as jnp
import optimistix as optx

y0 = jnp.array(1.)
dt = jnp.array(0.1)
def fn(y, args):
    return y0 + jnp.tanh(y) * dt

solver = optx.Newton(rtol=1e-5, atol=1e-5)
sol = optx.fixed_point(fn, solver, y0)
y1 = sol.value

Requires Python 3.11+; JAX and its dependencies (including jax itself) must be installed and functional.

Verify before relying

• Whether solver convergence guarantees or failure modes are documented for each algorithm.
• Performance comparison with other JAX optimization libraries or non-JAX solvers.
• Supported problem sizes and scalability limits on GPU/TPU.

Package facts

Evidence: optimistix-0.1.0-py3-none-any.whl

Keywords: deep-learning, equinox, jax, levenberg-marquardt, neural-networks, nonlinear-programming, numerical-optimization, optimization

Tags

More Artificial Intelligence packages