Lineax: unified linear solves and linear least-squares in JAX and Equinox
Jason Rader, Terry Lyons, Patrick Kidger

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.
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Code & Models
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsScientific Computing and Data Management
