Iterative linear solvers in the JAX ecosystem

Use this page when a linear system is too large or structured for a dense factorization and you need to interpret an iterative solver’s residual and derivative contract.

The scientific question¶

An iterative linear solver seeks an adequate approximation without explicitly forming or factorizing a large matrix. The scientific question is whether the operator structure, preconditioner, stopping rule, and numerical precision make the resulting state accurate enough for the downstream observable.

A small residual measures consistency with the represented linear system. It does not directly measure forward error when the operator is ill-conditioned, and it says nothing about whether the linearized model itself is scientifically appropriate.

Mathematical objects¶

Let $\mathbf{A}:\mathbb{R}^n\rightarrow\mathbb{R}^n$ be a matrix or linear operator, with right-hand side $\mathbf{b}$ and iterate $\mathbf{x}_k$ . A preconditioner approximates a useful inverse action while remaining cheaper to apply than solving the original system.

The norm, scaling, initial guess, finite precision, operator symmetry, positive definiteness, and preconditioner side all affect which algorithm and stopping interpretation are valid.

Core derivation¶

The residual generates a sequence of directions through repeated operator application. The associated Krylov space is:

\begin{aligned} \mathbf{A}\mathbf{x} &= \mathbf{b}, \\ \mathbf{r}_k &= \mathbf{b}-\mathbf{A}\mathbf{x}_k, \\ \mathcal{K}_m(\mathbf{A},\mathbf{r}_0) &=\mathrm{span}\left\{ \mathbf{r}_0,\mathbf{A}\mathbf{r}_0,\ldots, \mathbf{A}^{m-1}\mathbf{r}_0 \right\}. \end{aligned}

(1)

An iterative method chooses an update from (1) according to its operator assumptions and optimality criterion. If $\mathbf{e}_k=\mathbf{x}-\mathbf{x}_k$ , then $\mathbf{r}_k=\mathbf{A}\mathbf{e}_k$ . Therefore a small residual can coexist with a large forward error when inverse amplification is large.

What the ecosystem already owns¶

Lineax owns JAX-native linear-operator and linear-solve abstractions, while JAX provides transformation machinery for custom linear solves. Solver recurrences, operator dispatch, transposes, convergence statuses, and differentiation rules belong to those owners.

What Jaxstro may add¶

Jaxstro may later add a narrow adapter for unit and scale preparation, scientific provenance, or evidence reporting when a real downstream consumer needs it. No iterative-solver adapter exists today.

Such an adapter must preserve operator structure and external result statuses. It must not pretend that a reported residual is a forward-error certificate or that every solved system has a scientifically supported derivative.

Evidence required before implementation¶

Evidence would need:

dense reference comparisons on small systems with known structure;
residual and forward-error reports over controlled condition numbers;
exact-solution tests for symmetric, nonsymmetric, and rank-deficient cases within the claimed scope;
preconditioner tests that preserve the physical solution and report setup;
transpose and implicit-derivative checks against independent finite differences on smooth, converged cases;
nonconvergence, breakdown, and nonfinite failure cases; and
provenance for operator assumptions, tolerances, iterations, and status.

Claim boundary¶

This page neither selects one universal Krylov method nor reports Lineax or JAX performance. It documents the mathematical spine and the delegation boundary.

Connected foundations and methods¶

Start with Linear algebra as the language of change and Sensitivity, conditioning, and identifiability. Current small dense helpers are described in Linear algebra helpers, reusable operator structure in Linear operators, and Jacobian product mechanics in Autodiff products.