# Optimistix: modular optimisation in JAX and Equinox

arXiv: https://arxiv.org/abs/2402.09983 · HTML: https://mlsanity.com/paper/2402.09983/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 15 Feb 2024.

## Built on by (most cited first, 19 of 19)

- [Optimizing the Optimizer for Physics-Informed Neural Networks and Kolmogorov-Arnold Networks](https://mlsanity.com/paper/2501.16371/) (2025)
- [Representation Meets Optimization: Training PINNs and PIKANs for Gray-Box Discovery in Systems Pharmacology](https://mlsanity.com/paper/2504.07379/) (2025)
- [torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch](https://mlsanity.com/paper/2601.13994/) (2026)
- [SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients](https://mlsanity.com/paper/2603.08824/) (2026)
- [Self-Scaled Broyden Family of Quasi-Newton Methods in JAX](https://mlsanity.com/paper/2603.10599/) (2026)
- [Tackling Failure Modes of PINNs and PIKANs Using Conflict-Free Gradients](https://mlsanity.com/paper/2609.14841/) (2026)
- [Accelerating Natural Gradient Descent for PINNs with Randomized Numerical Linear Algebra](https://mlsanity.com/paper/2505.11638/) (2025)
- [Towards real-time surrogate-free Bayesian inversion for neutron reflectometry](https://mlsanity.com/paper/2509.06924/) (2025)
- [Automated design of compound lenses with discrete-continuous optimization](https://mlsanity.com/paper/2509.23572/) (2025)
- [Meshless solutions of PDE inverse problems on irregular geometries](https://mlsanity.com/paper/2510.25752/) (2025)
- [Case study of a differentiable heterogeneous multiphysics solver for a nuclear fusion application](https://mlsanity.com/paper/2511.13262/) (2025)
- [Second-Order, First-Class: A Composable Stack for Curvature-Aware Training](https://mlsanity.com/paper/2603.25976/) (2026)
- [Accelerated Fourier SAT (AFSAT): Fully Realising a GPU-based Symmetric Pseudo-Boolean SAT Solver](https://mlsanity.com/paper/2606.06641/) (2026)
- [RAM: Reachability Across Morphologies](https://mlsanity.com/paper/2606.09108/) (2026)
- [Differentiable design of the PIAA-ZWFS: a flexible wavefront sensor that approaches the fundamental limit](https://mlsanity.com/paper/2606.28136/) (2026)
- [SoRoMoX: Fast, Differentiable, and Parallelizable Soft Robot Models](https://mlsanity.com/paper/2608.06650/) (2026)
- [GRADSOLVE: fast exact gradients for ODE ensembles on GPUs](https://mlsanity.com/paper/2609.02876/) (2026)
- [Composite-Gradient Learning for Shared Control Authority Between Deep Reinforcement Learning and Model Predictive Control](https://mlsanity.com/paper/2609.17697/) (2026)
- [Coordinate-Free Automatic Differentiation](https://mlsanity.com/paper/2610.12776/) (2026)

## Abstract (arXiv)

We introduce Optimistix: a nonlinear optimisation library built in JAX and Equinox. Optimistix introduces a novel, modular approach for its minimisers and least-squares solvers. This modularity relies on new practical abstractions for optimisation which we call search and descent, and which generalise classical notions of line search, trust-region, and learning-rate algorithms. It provides high-level APIs and solvers for minimisation, nonlinear least-squares, root-finding, and fixed-point iteration. Optimistix is available at https://github.com/patrick-kidger/optimistix.