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

2023/11/28 by Jason Rader, Terry Lyons, Rader, Jason +3 · 4 citations
Decision Sciences · #FOS: Computer and information sciences #Mathematical Software (cs.MS) #Scientific Computing and Data Management

paper · pdf · doi:10.48550/arxiv.2311.17283

openalex publication_date 2023/11/28 · openalex created_date 2023/12/01 · openalex updated_date 2026/07/28

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.

Cited by

• jaxdae: A JAX-native Differentiable Solver for Differential-Algebraic Equations in Coupled Multi-physics
• JAX-AMG: A GPU-Accelerated Differentiable Sparse Linear Solver Library for JAX
• Differentiating Through a Quadratic Cone Program
• Structure-Preserving Discretization and Model Reduction for Energy-Based Models

Related

• MPAX: Mathematical Programming in JAX
• linrax: A JAX Compatible, Simplex Method Linear Program Solver
• Optimistix: modular optimisation in JAX and Equinox
• JAX-AMG: A GPU-Accelerated Differentiable Sparse Linear Solver Library for JAX
• JAXFit: Trust Region Method for Nonlinear Least-Squares Curve Fitting on the GPU