lineax: linear solves and least squares as JAX operators

Solving Ax = b when A is ill-posed or rectangular

Lineax is a JAX library for linear solves and linear least squares, and the scope statement in the documentation is precise: it provides routines that solve for x in Ax = b, and it keeps working when A may be ill-posed or rectangular. That last clause is the whole point of the package. Most numerical linear algebra libraries assume a square, well-conditioned matrix, and a large part of the practical difficulty in optimisation and inverse problems is that neither assumption holds.

The feature list gives a sense of the design. Matrices and vectors can be PyTree valued, so they need not be dense arrays and can carry arbitrary structure. There are general linear operators for Jacobians, transposes and similar constructions. Least squares can be solved efficiently, for instance with QR solvers. Gradients through a least squares solve are numerically stable, which is a separate problem from the solve itself and one that naive implementations get wrong. Structured matrices, such as symmetric ones, are supported. The documentation also claims improved compilation times and improved runtime of some algorithms, and both real and complex valued inputs are supported.

The last item on the list is the one that ties the library to its ecosystem. Working inside JAX means the solve is a JAX transformation, so it composes with autodiff, with jit compilation, with automatic parallelisation across devices, and with GPU and TPU execution. A solve that is a traced function is a solve that can live inside a training step.

Two ways to describe the operator, and one that never materialises a matrix

The quick examples show two ways to state the same problem, and the difference between them is the reason to use this library.

The first builds an operator from an explicit matrix and solves a least squares problem with a QR solver:

import jax.random as jr
import lineax as lx

matrix_key, vector_key = jr.split(jr.PRNGKey(0))
matrix = jr.normal(matrix_key, (10, 8))
vector = jr.normal(vector_key, (10,))
operator = lx.MatrixLinearOperator(matrix)
solution = lx.linear_solve(operator, vector, solver=lx.QR())

Note the shapes. The matrix is 10 by 8 and the vector has 10 entries, so the system is overdetermined, and the QR solver is the appropriate choice for that case. The random key is split with jr.split, which is the modern JAX idiom for deriving independent keys, and the operator wraps the matrix rather than the matrix being passed in directly.

The second example solves a quadratic problem without ever materialising a matrix:

import jax
import lineax as lx

key = jax.random.PRNGKey(0)
y = jax.random.normal(key, (10,))

def quadratic_fn(y, args):
  return jax.numpy.sum((y - 1)**2)

gradient_fn = jax.grad(quadratic_fn)
hessian = lx.JacobianLinearOperator(gradient_fn, y, tags=lx.positive_semidefinite_tag)
solver = lx.CG(rtol=1e-6, atol=1e-6)
out = lx.linear_solve(hessian, gradient_fn(y, args=None), solver)
minimum = y - out.value

Here the operator is a JacobianLinearOperator built from the gradient of the quadratic, which means the Hessian is defined by autodiff and never stored. The tag argument, positive_semidefinite_tag, tells the solver a property of the operator that it cannot infer, and the solver is conjugate gradient with a relative and an absolute tolerance of 1e-6. The result is read from out.value, and the minimum is recovered by stepping back along the gradient.

Operators, tags and solvers are the three decisions you make

Read the two examples together and the API surface is clear. A linear operator describes the system, a solver decides how it is solved, and tags tell the solver about structure it cannot see.

The operators are where the library earns its keep. A MatrixLinearOperator wraps something you already hold. A JacobianLinearOperator is defined by a function, in this case the gradient of a scalar function, and applying it means evaluating that function and applying its Jacobian, without forming the matrix. For a model with millions of parameters, the difference between forming a Hessian and never forming it is the difference between a method that runs and one that does not.

Tags are the third piece, and they are the least obvious. The tag argument on the operator is not documentation, it is information passed to the solver about the structure of the operator, in this case that it is positive semidefinite. A solver that knows the matrix is positive semidefinite, or symmetric, or complex can choose a method and a factorisation that are valid under that assumption. Getting the tag right is what lets the solver be both faster and more accurate than a general one, and getting it wrong is how a solver silently returns the wrong answer for a problem that is not what you claimed it was.

The solver choice follows the operator. QR is the right answer for an overdetermined dense system, as in the first example. Conjugate gradient is the right answer for a large positive semidefinite operator where the matrix is only available through its action, as in the second, and it comes with explicit tolerances, rtol and atol, set to 1e-6 in the example. The result object carries the solution and, presumably, the diagnostics you need to decide whether the solve actually converged, since a solver that stops at its tolerance without saying so is worse than no solver at all.

Installing it, and the version requirements that disagree with themselves

Installation is a single command:

pip install lineax

The requirements that follow it are where the documentation and the package metadata part company. The documentation states that Lineax requires Python 3.10+, JAX 0.4.38+ and Equinox 0.11.10 or newer. The package metadata in pyproject.toml says requires-python = "~=3.11", which means 3.11 and compatible releases only, and lists the runtime dependencies as jax>=0.10.0, jaxtyping>=0.2.24, equinox>=0.11.10 and typing_extensions>=4.5.0.

So the two sources disagree on two points. The documented Python floor is looser than the one the package actually declares, which means pip will refuse an install on Python 3.10 even though the documentation says it is supported. And the documented JAX floor of 0.4.38 is a different kind of statement from the declared jax>=0.10.0, since a lower bound that has been overtaken by later releases no longer tells you which JAX version was tested. If your environment is pinned to an older JAX, that is the question to resolve before you install, and the answer is not in the README.

The tooling around the package is modern and conventional. The build backend is hatchling, and the project metadata names Jason Rader and Patrick Kidger as the authors, classifies the project as Development Status :: 3 - Alpha, and declares the licence as Apache Software License from the LICENSE file. The dev dependency group is pinned exactly, with prek, pyright, ruff and toml-sort, the docs group brings in mkdocs with its Material theme and mkdocstrings, and the tests group lists beartype, equinox, pytest, pytest-xdist and jaxlib, which tells you the test suite runs in parallel and checks runtime types. The repository layout follows: lineax/ for the package, tests/, benchmarks/, docs/ with an mkdocs.yml at the top, and a .pre-commit-config.yaml for the formatting and lint hooks.

Where it sits in the JAX ecosystem, and the paper behind it

The project documentation ends with a map of the surrounding ecosystem, which is useful because Lineax is one piece of a set rather than a standalone tool.

The always useful entries are Equinox, for neural networks and everything not already in core JAX, and jaxtyping, for type annotations on the shape and dtype of arrays. The deep learning group lists Optax for first-order gradient optimisers, Orbax for checkpointing across hosts and devices, Levanter for scalable and reliable training of foundation models, and paramax for parameterizations and constraints on PyTrees. The scientific computing group lists Diffrax for numerical differential equation solvers, Optimistix for root finding, minimisation, fixed points and least squares, BlackJAX for probabilistic and Bayesian sampling, sympy2jax for converting between SymPy and JAX, and PySR for symbolic regression.

That list is also the honest alternative to Lineax in some cases. Optimistix covers least squares as part of a broader minimisation package, so if your problem is minimisation rather than a linear solve, the same ecosystem already contains an answer, and reaching for Lineax may mean implementing a solver you could have configured. The distinction is that Lineax is for the linear algebra itself, and for the cases where the operator is implicit.

The academic reference is a paper, cited with a BibTeX entry: Lineax, unified linear solves and linear least-squares in JAX and Equinox, by Jason Rader, Terry Lyons and Patrick Kidger, from the AI for science workshop at Neural Information Processing Systems 2023, with the arXiv identifier 2311.17283. For a library whose design choices are unusual, the paper is the right place to look for the reasoning behind them, and the citation request in the documentation is the standard academic courtesy of the author.

Alpha status, JAX lock-in and the questions to ask first

Three limits are worth stating before anyone builds on this.

The first is maturity. The project classifies itself as alpha, and the release history supports caution rather than alarm: version 0.1.0 on 2026-01-27, version 0.1.1 on 2026-05-01, and version 0.0.8 on 2025-03-27 before that. The repository is not archived and the last push was on 2026-09-26, so the work continues, and the jump from 0.0.x to 0.1.x with a patch release inside four months is a normal early rhythm. The API can still move.

The second is the JAX commitment. There is no NumPy path, and a solver that is a traced JAX function cannot be dropped into a numerical code base that has not already accepted that dependency. If your linear algebra is currently in SciPy, adopting Lineax means adopting JAX, and the change is bigger than the import.

The third is the numerical responsibility. The library gives you solvers, tolerances and a structured description of the operator, and it will happily solve an ill-posed system, which means the answer may be a least-norm solution that satisfies the equations without being the answer you wanted. Choosing between QR, conjugate gradient and the other solvers, and setting the structure tags correctly, is part of using it well, and the documentation for solvers is where you should read before picking one.

The concrete checks to run first are these. Resolve the Python and JAX version question against your environment, since the README and the package metadata disagree. Read the solver documentation and pick deliberately rather than by habit. If your matrices are dense and small, use an existing array library and skip the dependency, and consider Lineax when the operator is implicit, which is the case it was built for.

Editorial conclusion

Lineax fits a research or production JAX codebase that repeatedly has to solve a linear system, and especially one where the matrix is implicit, such as a Hessian obtained from autodiff, where a conventional solver would first have to materialise something that is expensive or impossible to hold. It is the wrong tool for a NumPy codebase, since everything in it depends on JAX transformations, and it is the wrong tool for a project that needs a stable, supported version, because the project classifies itself as alpha. Before you adopt it, check the version requirements against your environment, because the documentation and the package metadata state different minimums for Python and JAX, and pick the solver deliberately, since the choice of QR against CG and the structure tags on the operator change the numerical result and not just the speed.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. patrick-kidger/lineax on GitHub
4. README
5. Releases