Lineax: unified linear solves and linear least-squares in JAX and Equinox

TL;DR

Lineax is a new library that unifies linear solves and least-squares problems in JAX+Equinox, offering an extensible, autodifferentiable API without needing custom derivatives.

Contribution

It introduces a unified, extensible API for linear solves and least-squares in JAX+Equinox, simplifying differentiation and customization.

Findings

Provides a single API for linear solves and least-squares

Ensures autodifferentiability without custom derivatives

Open-source implementation available at GitHub

Abstract

We introduce Lineax, a library bringing linear solves and linear least-squares to the JAX+Equinox scientific computing ecosystem. Lineax uses general linear operators, and unifies linear solves and least-squares into a single, autodifferentiable API. Solvers and operators are user-extensible, without requiring the user to implement any custom derivative rules to get differentiability. Lineax is available at https://github.com/google/lineax.