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.
@article{arxiv.2311.17283,
title = {Lineax: unified linear solves and linear least-squares in JAX and Equinox},
author = {Jason Rader and Terry Lyons and Patrick Kidger},
journal= {arXiv preprint arXiv:2311.17283},
year = {2023}
}
Comments
7 pages, 1 figure, NeurIPS 2023 AI for Science workshop