SOLVAX¶
Differentiable structured solvers, preconditioners, and matrix-free methods for JAX.
SOLVAX is the numerical-solver layer that kinetic, transport, equilibrium, and
PDE codes often reimplement independently. It provides structured direct
factorizations, matrix-free Krylov methods, fixed-point acceleration,
preconditioners, mixed-precision refinement, bounded-memory Jacobians, and
implicit differentiation. With the exception of the explicitly host-side
SuperLU bridge, the library is designed to compose with jax.jit, jax.vmap,
and jax.grad.
This documentation is organized around decisions rather than modules:
Getting started establishes the operator, tolerance, shape, and result conventions used everywhere.
Choosing a method maps mathematical structure to a solver and preconditioner.
The solver guides derive each algorithm, document every input and output, describe failure modes, and compare the method with common alternatives.
The tutorials build complete structured, matrix-free, and differentiable workflows.
API reference is the generated signature-level reference.
A first solve¶
import jax.numpy as jnp
import solvax as sx
A = jnp.array([[4.0, 1.0], [1.0, 3.0]])
b = jnp.array([1.0, 2.0])
solution = sx.pcg(lambda x: A @ x, b, rtol=1e-10)
if not solution.converged:
raise RuntimeError(sx.status_name(solution.status))
print(solution.x)
print(solution.residual_norm)
For a nonsymmetric operator, use FGMRES:
solution = sx.gmres(lambda x: A @ x, b, restart=20, rtol=1e-10)
For a block-tridiagonal operator, do not discard the structure:
factors = sx.block_thomas_factor(lower, diagonal, upper)
x = sx.block_thomas_solve(factors, rhs)
Documentation map¶
Structured direct solvers
Iterative and nonlinear solvers
Solver infrastructure
- Linear operator models
- Preconditioners
- Jacobi scaling
- Block Jacobi
- Coarse or simplified operator
- Line smoother
- Symmetric additive composition
- Multigrid V-cycle
- Randomized Nystrom preconditioning
- Symmetric Galerkin deflation
- Kronecker preconditioning
- Mixed-precision wrapper
- Constraint-aware preconditioning
- How to evaluate a preconditioner
- Compatibility table
- API summary
- Matrix-free nonlinear solves and implicit differentiation
- Mixed-precision iterative refinement
- Memory-bounded Jacobians
- Sharding and communication
- Native SuperLU bridge
Tutorials
Benchmarks & testing
Literature¶
SOLVAX follows established numerical linear algebra rather than introducing new convergence theory. Each method page states the algorithmic variant used by the implementation and cites the relevant literature. The full bibliography is collected below.
Donald G. Anderson. Iterative procedures for nonlinear integral equations. Journal of the ACM, 12:547–560, 1965. doi:10.1145/321296.321305.
Michele Benzi, Paola Boito, and Nader Razouk. Decay properties of spectral projectors with applications to electronic structure. SIAM Review, 55(1):3–64, 2013. doi:10.1137/100814019.
Michele Benzi, Gene H. Golub, and Jörg Liesen. Numerical solution of saddle point problems. Acta Numerica, 14:1–137, 2005. doi:10.1017/S0962492904000212.
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. In Advances in Neural Information Processing Systems. 2022. URL: https://arxiv.org/abs/2105.15183.
Erin Carson and Nicholas J. Higham. Accelerating the solution of linear systems by iterative refinement in three precisions. SIAM Journal on Scientific Computing, 40:A817–A847, 2018. doi:10.1137/17M1140819.
Stephen Demko, William F. Moss, and Philip W. Smith. Decay rates for inverses of band matrices. Mathematics of Computation, 43(168):491–499, 1984. doi:10.1090/S0025-5718-1984-0758197-9.
James W. Demmel, Nicholas J. Higham, and Robert S. Schreiber. Stability of block LU factorization. Numerical Linear Algebra with Applications, 2:173–190, 1995. doi:10.1002/nla.1680020208.
F. Javier Escoto. Fast and accurate calculation of the bootstrap current and radial neoclassical transport in low collisionality stellarator plasmas. PhD thesis, Universidad Carlos III de Madrid, 2025. URL: https://arxiv.org/abs/2510.27513.
Zachary Frangella, Joel A. Tropp, and Madeleine Udell. Randomized Nyström preconditioning. SIAM Journal on Matrix Analysis and Applications, 44:718–752, 2023. doi:10.1137/21M1466244.
Luc Giraud, Julien Langou, Miroslav Rozloznik, and Jasper van den Eshof. Rounding error analysis of the classical Gram-Schmidt orthogonalization process. Numerische Mathematik, 101:87–100, 2005. doi:10.1007/s00211-005-0615-4.
Gene H. Golub and Charles F. Van Loan. Matrix Computations. Johns Hopkins University Press, 4 edition, 2013.
Magnus R. Hestenes and Eduard Stiefel. Methods of conjugate gradients for solving linear systems. Journal of Research of the National Bureau of Standards, 49:409–436, 1952. doi:10.6028/jres.049.044.
Nicholas J. Higham. Accuracy and Stability of Numerical Algorithms. SIAM, 2 edition, 2002. doi:10.1137/1.9780898718027.
Dana A. Knoll and David E. Keyes. Jacobian-free Newton–Krylov methods: a survey of approaches and applications. Journal of Computational Physics, 193:357–397, 2004. doi:10.1016/j.jcp.2003.08.010.
Ronald B. Morgan. GMRES with deflated restarting. SIAM Journal on Scientific Computing, 24:20–37, 2002. doi:10.1137/S1064827599364659.
Christopher C. Paige and Michael A. Saunders. Solution of sparse indefinite systems of linear equations. SIAM Journal on Numerical Analysis, 12:617–629, 1975. doi:10.1137/0712047.
Michael L. Parks, Eric de Sturler, Greg Mackey, Duane D. Johnson, and Spandan Maiti. Recycling Krylov subspaces for sequences of linear systems. SIAM Journal on Scientific Computing, 28:1651–1674, 2006. doi:10.1137/040607277.
William H. Press, Saul A. Teukolsky, William T. Vetterling, and Brian P. Flannery. Numerical Recipes: The Art of Scientific Computing. Cambridge University Press, 3 edition, 2007. Section 2.7, cyclic tridiagonal systems.
Yousef Saad. Iterative Methods for Sparse Linear Systems. SIAM, 2 edition, 2003. doi:10.1137/1.9780898718003.
Calum S. Skene and Keaton J. Burns. Fast automated adjoints for spectral PDE solvers. Proceedings of the National Academy of Sciences, 2026. URL: https://arxiv.org/abs/2506.14792.
Paul N. Swarztrauber. The methods of cyclic reduction, Fourier analysis and the FACR algorithm for the discrete solution of Poisson's equation on a rectangle. SIAM Review, 19:490–501, 1977. doi:10.1137/1019071.
Llewellyn H. Thomas. Elliptic problems in linear difference equations over a network. Technical Report, Watson Scientific Computing Laboratory, Columbia University, 1949.
Ulrich Trottenberg, Cornelis W. Oosterlee, and Anton Schüller. Multigrid. Academic Press, 2001.
Henk A. van der Vorst. Bi-CGSTAB: a fast and smoothly converging variant of Bi-CG for the solution of nonsymmetric linear systems. SIAM Journal on Scientific and Statistical Computing, 13:631–644, 1992. doi:10.1137/0913035.
Charles F. Van Loan and Nikos Pitsianis. Approximation with Kronecker products. In Linear Algebra for Large Scale and Real-Time Applications. Springer, 1993. doi:10.1007/978-94-015-8196-7_17.
Homer F. Walker and Peng Ni. Anderson acceleration for fixed-point iterations. SIAM Journal on Numerical Analysis, 49:1715–1735, 2011. doi:10.1137/10078356X.