Lineax: Unified Linear Solves and Linear Least-Squares in JAX and 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 https://github.com/$\textbf{anonymised}$/lineax.

PDF NeurIPSW OpenReview Semantic Scholar

Cite

Text

Rader et al. "Lineax: Unified Linear Solves and Linear Least-Squares in JAX and Equinox." NeurIPS 2023 Workshops: AI4Science, 2023.

Markdown

[Rader et al. "Lineax: Unified Linear Solves and Linear Least-Squares in JAX and Equinox." NeurIPS 2023 Workshops: AI4Science, 2023.](https://mlanthology.org/neuripsw/2023/rader2023neuripsw-lineax/)

BibTeX

@inproceedings{rader2023neuripsw-lineax,
  title     = {{Lineax: Unified Linear Solves and Linear Least-Squares in JAX and Equinox}},
  author    = {Rader, Jason Michael and Lyons, Terry and Kidger, Patrick},
  booktitle = {NeurIPS 2023 Workshops: AI4Science},
  year      = {2023},
  url       = {https://mlanthology.org/neuripsw/2023/rader2023neuripsw-lineax/}
}