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

Published: 28 Oct 2023, Last Modified: 28 Nov 2023NeurIPS2023-AI4Science PosterEveryoneRevisionsBibTeX
Keywords: Jax, Equinox, Open source, Scientific computing, Scientific machine learning, Numerical linear algebra
TL;DR: We introduce Lineax: an open-source library for unified, autodifferentiable linear solves and least-squares in JAX+Equinox.
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$\textbf{anonymised}$/lineax.
Submission Track: Original Research
Submission Number: 68