SOLVAX

Differentiable structured linear solvers, preconditioners and matrix-free methods in JAX.

solvax provides the solver infrastructure that kinetic and PDE codes keep re-implementing: structured direct solves (batched dense LU, block-tridiagonal Schur elimination with truncated storage), preconditioned and recycled Krylov methods, physics-agnostic preconditioners (coarse-operator LU, semicoarsened geometric and p-multigrid, Kronecker approximations, symmetric additive and line smoothers), mixed-precision iterative refinement, implicit differentiation wrappers for converged linear and nonlinear solves, checkpointed long recurrences, and memory-chunked Jacobians — traceable under jit, vmap and grad, on CPU and GPU.

Derivative semantics are explicit. Use pcg_linear_solve, linear_solve, affine_fixed_point_gmres, or root_solve for a converged-equation derivative that does not tape iterations. Raw PCG/GMRES/GCROT, relaxed/Aitken iterations, and Newton--Krylov are primal algorithms. Direct and truncated structured solvers document whether their derivative is taped, implicit, exact-window, or reverse-only.

Three documented exceptions, because "transparent to every transform" would not be true: the exact-window reverse rule is a custom_vjp, so jax.jacfwd and jax.jvp raise on it rather than falling back to the taped path (pass the full window, or differentiate the untruncated entry point, when forward mode is what you need); eigenpair_reverse supports reverse mode but orchestrates application-supplied eigensolvers on the host, so their compiled kernels sit inside rather than around that dispatcher; and solvax.native host primitives refuse direct tracing while still composing as external primals with eigenpair_reverse.

It complements general JAX solver libraries with block-structured direct elimination, coarse-operator and multigrid preconditioning, and Krylov subspace recycling for parameter continuation. SOLVAX operators are native JAX pytrees; no external operator abstraction is required.

Install

pip install solvax

Quickstart

import jax
import jax.numpy as jnp
import solvax as sx

# Solve a block-tridiagonal system L_k x_{k-1} + D_k x_k + U_k x_{k+1} = b_k
x = sx.block_thomas(lower, diag, upper, rhs)

# Matrix-free PCG on arrays or arbitrary JAX pytrees
solution = sx.pcg(matvec, rhs, precond=preconditioner, rtol=1e-10)
assert solution.converged

# Solve an expensive affine coupling map without assembling its Jacobian
coupled = sx.affine_fixed_point_gmres(coupling_sweep, initial_state)

# Run a relaxed nonlinear coupling sweep with explicit stopping diagnostics
fixed = sx.fixed_point_iteration(
    coupling_sweep, initial_state, relaxation=0.8, atol=1e-8
)

# Periodic Poisson solve; reuse the symbol when the timestep owns the FFT
grid = 2.0 * jnp.pi * jnp.arange(16) / 16
rho = jnp.sin(grid[:, None]) * jnp.sin(grid[None, :])
dx = dy = 2.0 * jnp.pi / 16
rho_hat = jnp.fft.fftn(rho)
phi = sx.solve_periodic_poisson(rho, spacing=(dx, dy))
poisson_symbol = sx.periodic_poisson_eigenvalues(rho.shape, (dx, dy))
phi_hat = sx.solve_periodic_poisson_spectral(rho_hat, eigenvalues=poisson_symbol)

# Same diagnostics, but gradients use an implicit primal/transpose solve
implicit_solution = sx.pcg_linear_solve(matvec, rhs, precond=preconditioner)

# Exact discrete gradient of a long recurrence with O(sqrt(steps)) retained state
def advance(_index, state):
    return coupling_sweep(state)


final_state = sx.checkpointed_fori_loop(0, 8, advance, initial_state)

# Reuse one elimination across many right-hand sides
factors = sx.block_thomas_factor(lower, diag, upper)
x1 = sx.block_thomas_solve(factors, rhs1)
x2 = sx.block_thomas_solve(factors, rhs2)

# Generate each block once when reusable factors are needed without a stored
# diagonal band. `row(j)` returns the triple (L_j, D_j, U_j) of row j; the
# parameterized form `block_fn(params, j)` used further down takes the
# parameters as its first argument.
generated_factors = sx.block_thomas_factor_fn(row, n_blocks=N)

# Same factors, a third of the state: keep the Schur LU only and rebuild the
# off-diagonal blocks from `row` during each substitution sweep.
lean_factors = sx.block_thomas_factor_fn(row, n_blocks=N, store_offdiagonals=False)
x_lean = sx.block_thomas_solve(lean_factors, rhs)
adjoint = sx.block_thomas_solve(lean_factors, rhs, transpose=True)

# One generated solve with O(sqrt(N) m^2) factor storage and exact JVP/VJP.
x = sx.block_thomas_checkpointed_fn(row, N, rhs)

# Memory-truncated mode: rhs nonzero only in the lowest K blocks and only the
# lowest K solution blocks needed -> O(K m^2) memory, independent of N.
x_low = sx.block_thomas_truncated(lower, diag, upper, rhs[:3], keep_lowest=3)

Differentiate a generated selected-head solve with respect to the compact parameters that build its rows, at retained state independent of the block count. Ask for the accuracy you want and get a window that provably delivers it:

# Smallest window whose relative gradient error is provably below rtol.
cert = sx.certified_adjoint_window(block_fn, N, keep_lowest=3, params=p,
                                   rhs_low=rhs[:3], cotangent=ct, rtol=1e-8)
print(cert.window, cert.certified_relative_error)  # proven, not estimated

def objective(params):
    x_low = sx.block_thomas_truncated_fn(
        block_fn, N, rhs[:3], keep_lowest=3,
        params=params, adjoint_window=cert,
    )
    return loss(x_low)

grad = jax.grad(objective)(p)

The certificate is possible because the exact-window rule leaves exactly one approximation to bound — the omitted rows — while every retained row cotangent is an exact block of the full solve. It costs two localization sweeps, one generator Jacobian per row, and a differentiated solve, so it is a setup computation rather than something to run each optimizer step.

Two cheaper options remain for when that is too much. localization_crossover_window reads the chain's transfer norms and returns a starting window with certified=False, and check_localized_gradient widens a window to see whether the gradient moves — evidence rather than a bound, but it needs no Jacobians.

For a batch of chains that localize at different rows — a collisionality scan, say — adjoint_window is static and would impose the worst chain's width on all of them. plan_chain_windows cuts the batch into a few static buckets instead, and chain_window= gives each chain its own window inside one:

plan = sx.plan_chain_windows(windows, keep_lowest=3, max_buckets=4)
print(plan.reduction)  # measured: 51% fewer retained rows on a 24-chain nu scan
for window, chains in plan.buckets:
    grads = jax.vmap(lambda c: chain_gradient(c, adjoint_window=window))(chains)

Everything is differentiable (jax.grad through the solve) and batchable (jax.vmap over stacked systems).

Differentiate a certified eigenmode without taping its solver

Applications often need a mode shape as well as its eigenvalue, but recording hundreds of Krylov or timestepper iterations makes reverse-mode optimization slow and memory hungry. eigenpair_reverse leaves mode finding to the application and differentiates only the converged simple eigenpair:

initial_mode = jnp.ones(2, dtype=jnp.complex128)
parameters = jnp.asarray(0.0)

def build(parameters):
    matrix = jnp.asarray(
        [[0.5 + 0.3 * parameters + 0.2j, 0.0],
         [parameters, -1.0 + 3.0j]],
        dtype=jnp.complex128,
    )
    return lambda vector: matrix @ vector

def primal_solver(parameters, _apply, _start):
    # A real application injects its residual-certified matrix-free solver.
    eigenvalue = 0.5 + 0.3 * parameters + 0.2j
    eigenvector = jnp.asarray(
        [1.0, parameters / (eigenvalue + 1.0 - 3.0j)],
        dtype=jnp.complex128,
    )
    return eigenvalue, eigenvector

def left_solver(_parameters, _apply, _start, _eigenvalue):
    return jnp.asarray([1.0, 0.0], dtype=jnp.complex128)

def objective(parameters):
    eigenvalue, eigenvector = sx.eigenpair_reverse(
        parameters,
        build,
        initial_mode,
        primal_solver=primal_solver,
        left_solver=left_solver,
    )
    normalized = eigenvector / jnp.linalg.norm(eigenvector)
    return jnp.real(eigenvalue) + 0.1 * jnp.real(normalized[1])

gradient = jax.grad(objective)(parameters)

For leftᴴ right = 1, the eigenvalue rule is dλ = leftᴴ (dA) right; eigenvector observables use one bordered reduced-resolvent solve. This is independent of state layout and primal algorithm, avoids differentiating the eigensolver iteration, and rejects nearly defective eigenpairs through an explicit condition-number gate. That gate reads the pair you supply, so it is bounded by your own solver's accuracy: near an exceptional point the solver degrades first, and the vectors it returns can make the condition number look acceptable when it is not. Treat a value within an order of magnitude of the limit as unconfirmed.

What's in the box

Complex-valued GMRES/GCROT, tridiagonal solves, and fixed-point acceleration use Hermitian inner products and real-valued safeguards. Remaining roadmap: multi-leaf pytree GCROT operands (GCROT takes arrays of any rank; GMRES is pytree-native) and expanded GPU batched-LU benchmarks.

# Preconditioned, recycled Krylov across a parameter scan:
sol = sx.gcrot(matvec, b, precond=coarse_inverse, m=50, k=10)
sol2 = sx.gcrot(matvec2, b2, precond=coarse_inverse, recycle=sol.recycle)

# Matrix-free Newton-Krylov (JFNK): Jacobian-vector products via jax.linearize,
# each correction solved by FGMRES over an array or structured pytree state.
root = sx.newton_krylov(residual_fn, x0, precond=approx_inverse, rtol=1e-8)

# Globalized root continuation with a pseudo-time metric and hard bounds.
root = sx.pseudo_transient_continuation(
    residual_fn, x0, mass=lambda state, value: value,
    admissible=lambda x: jnp.all(jnp.isfinite(x)),
)

# Weakly contractive affine coupling map G(x) = L x + c, solved as (I - L) x = c:
fixed = sx.affine_fixed_point_gmres(coupling_map, x0, restart=20)

# Periodic (cyclic) scalar tridiagonal line, corners in sub[0] and sup[-1]:
x_line = sx.cyclic_tridiagonal_solve(sub, dia, sup, line_rhs)

# Differentiable solve wrapping any solver:
x = sx.linear_solve(matvec, b, solver=lambda mv, rhs: sx.gmres(mv, rhs).x)

License

MIT. Developed by the UW Plasma group.