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, and implicit differentiation of every solve — all jit/vmap/grad-transparent, on CPU and GPU.

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.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)

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

# 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.
generated_factors = sx.block_thomas_factor_fn(block_fn, n_blocks=N)

# One generated solve with O(sqrt(N) m^2) factor storage and exact JVP/VJP.
x = sx.block_thomas_checkpointed_fn(block_fn, 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. The window is chosen up front from the chain's own localization profile:

w = sx.suggest_adjoint_window(lambda k: block_fn(p, k), N, keep_lowest=3)
x_low = sx.block_thomas_truncated_fn(
    block_fn, N, rhs[:3], keep_lowest=3, params=p, adjoint_window=w
)
grad = jax.grad(lambda p: loss(x_low))(p)

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

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)

# 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) tridiagonal line, corners in lower[0] and upper[-1]:
x = sx.cyclic_tridiagonal_solve(lower, diag, upper, 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.