Optimistix is a JAX library for solving nonlinear equations, targeting researchers and engineers working with JAX and Equinox. It offers a modular and interoperable approach to root-finding, minimization, fixed-point problems, and least squares, leveraging JAX's automatic differentiation, parallelization, and hardware acceleration.

How It Works

Optimistix provides a highly modular design, allowing users to combine different solvers, descent methods, and trust region updates. For instance, a root-finding problem can be automatically converted to a least-squares problem and solved with a minimization algorithm. This flexibility, coupled with its PyTree-based state management and seamless integration with Optax, enables fast compilation and runtime performance.

Quick Start & Requirements

• Install via pip: pip install optimistix
• Requires Python 3.10+, JAX 0.4.38+, and Equinox 0.11.11+.
• Documentation: https://docs.kidger.site/optimistix/

Highlighted Details

• Solvers can be interoperably combined (e.g., root-finding to least-squares to minimization).
• Modular optimizers allow mixing descent paths and trust region updates.
• Leverages JAX for autodiff, autoparallelism, and GPU/TPU support.
• Integrates with Optax for first-order optimization.

Maintenance & Community

Primarily built by Jason Rader. The project is part of a broader JAX ecosystem, with related libraries like Equinox, Diffrax, and Lineax.

Licensing & Compatibility

The README does not explicitly state a license. Compatibility for commercial use or closed-source linking is not specified.

Limitations & Caveats

The library is primarily focused on JAX and Equinox users. The absence of an explicit license in the README may require further investigation for commercial adoption.