Optimistix: modular optimisation in JAX and Equinox

TL;DR

Optimistix is a new modular nonlinear optimisation library built on JAX and Equinox, offering flexible abstractions for various optimisation algorithms and high-level APIs for different problem types.

Contribution

It introduces a novel modular framework with abstractions called search and descent, generalising classical optimisation methods.

Findings

Provides a flexible, modular optimisation library in JAX and Equinox.

Offers high-level APIs for minimisation, root-finding, and fixed-point iteration.

Enables easy customization and extension of optimisation algorithms.

Abstract

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.