arXiv Science⌕ Search
Home/Archive/arXiv

arXiv · 2606.09001

JAX-AMG: A GPU-Accelerated Differentiable Sparse Linear Solver Library for JAX

Yi Liu·Xiantao Fan·Jian-Xun Wang
View original source ↗Read source document ↗Export metadata ↓

Abstract

Sparse linear systems from PDE discretizations are central to scientific computing, yet no existing JAX-ecosystem solver simultaneously provides GPU-accelerated algebraic multigrid (AMG), automatic differentiation (AD), and distributed multi-GPU execution. JAX-AMG fills this gap by wrapping the NVIDIA AmgX solver suite as a native JAX primitive, exposing AMG and Krylov methods with configurable preconditioners through a unified interface compatible with JIT compilation, reverse-mode AD via adjoint methods, batched solves, and MPI-based distributed execution. Solver caching amortizes setup costs across repeated solves, making JAX-AMG practical for PDE-constrained optimization and inverse problems. The result is a robust, scalable sparse linear algebra layer that integrates seamlessly into differentiable simulation and scientific machine learning pipelines.

Explore related subjects

cs.MSphysics.comp-ph

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Yi Liu, Xiantao Fan, Jian-Xun Wang. 2026-07-28. JAX-AMG: A GPU-Accelerated Differentiable Sparse Linear Solver Library for JAX. https://doi.org/10.1016/j.softx.2026.102966

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Eigensolvers for polynomial roots and tensor decomposition

Computing eigenvalues and eigenvectors is at the heart of the solution of many non-linear problems. For instance, finding the roots of polynomial systems reduces to computing joint eigenvectors of operators of multiplication. Similarly, tensor decomposition can be performed via the joint diagonalization of submatrices of the Catalecticant of the tensor. We describe and illustrate symbolic-numeric methods for computing the solutions of these algebraic problems from the computation of joint eigenvectors of commuting operators, and for analysing their multiplicity structure, as well as their implementation in the package AlgebraicSolvers.jl.

cs.MS↗
arXiv2026

The Art of Closed-Formula Defaults: Search-Free Code Generation for Tensor Operators

Agentic search and automated optimization of GPU kernels are powerful tools for large language model inference. Their effectiveness, however, depends not on the sophistication of the search itself, but on the clarity of the optimization problem being solved. We provide an application-first approach that drives a hierarchical code generation tool from operator specifi cation down to GPU instructions, and show that a clearly defined computational model makes the optimization problem tractable.

cs.MS↗
arXiv2026

FP8 is All You Need (Part 2): Full-FP64 3-D FFT on FP8-Generation Tensor CoresThe Integer-Epilogue Wall and the Minimal Hardware That Would Remove It

The NVIDIA Blackwell Ultra (B300) GPU cuts FP64 vector throughput $\sim 30\times$ while multiplying FP8 tensor throughput. After the recovery of FP64 GEMM via Ozaki Scheme II on FP8 tensor cores and the Tensor-Memory Equilibrium model of the companions ("FP8 is All You Need, Part 1" and "Ozaki 2.5") we ask whether the fifth canonical HPC primitive, the full-FP64 $1024^3$ 3-D FFT, can be carried by the same substrate, and answer with a design and its limit. It is a Bailey six-step transform with no FP64 arithmetic: FP8-tensor DFT GEMMs with fused twiddles, residue-domain Karatsuba combines and exact CRT reconstruction whose bulk is a small GEMM on the FP16 tensor path and whose remainder is a Kulisch fixed-point accumulation with a two-sided modulo-$M$ lift, so the only rounding is the final conversion; constants are machine-generated and verified bit-exactly. The central finding: the binding resource is not floating point but a per-output integer epilogue with floor $(c_{\rm epi}/8),B_{\rm mem}$, $c_{\rm epi} \approx 203$-$281$ instructions per output: on B300 it holds the transform at 63-87 ms against a 12.9 ms roof ($4.9$-$6.7\times$ short); at most $1.3$-$1.9\times$ faster than the collapsed native path, possibly no faster at realised issue rates; no software route reaches the roof; on the NVIDIA Rubin GPU emulation loses $8$-$11\times$. An FP32 variant meets the same wall: the cause is per-scalar reconstruction, not FP64. Each floor term names its remedy: the NVIDIA B200 GPU's INT8 tensor core restored with a position-weighted cross-column accumulation primitive, a load-path deconstruction datapath shared with the companions, two ISA idioms and modular reduction at the MMA output give 16.0-23.5 ms with minor hardware and 12.9-15.0 ms with one moderate ask. All figures are projected floors, not measurements, with sensitivities and the FP8 layout condition given.

cs.MS↗