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

Yi Liu, Xiantao Fan, Jian-Xun Wang

Sparse linear systems from PDE discretizations are central to scientific computing, yet no existing JAX-ecosystem solver simultaneously provides GPU-accelerated algebraic multigrid…