JAX-AMG: A GPU-accelerated differentiable sparse linear solver library for JAX

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

Metadata

Code metadata.

Nr.	Code metadata description	Metadata
C1	Current code version	v0.1.4
C2	Permanent link to code/repository used for this code version	https://github.com/jx-wang-s-group/JAX-AMG/tree/v0.1.4
C3	Permanent link to Reproducible Capsule
C4	Legal Code License	Apache License 2.0
C5	Code versioning system used	git
C6	Software code languages, tools, and services used	Python, C++, CUDA, MPI
C7	Compilation requirements, operating environments & dependencies	Python 3.10+, CUDA Toolkit 12+, NVIDIA AmgX 2.5+, JAX 0.5.0+ (with CUDA support), MPI library (with mpi4py and mpi4jax)
C8	If available Link to developer documentation/manual	https://jx-wang-s-group.github.io/JAX-AMG/
C9	Support email for questions	jw2837@cornell.edu

1. Motivation and significance

Sparse linear systems arising from the discretization of partial differential equations (PDEs) are ubiquitous in computational science and engineering, with critical applications across fluid mechanics [1], heat transfer [2], structural mechanics [3], electromagnetics [4], and many other disciplines. Increasingly, the solvers for these systems are embedded within gradient-based optimization and machine learning pipelines that require end-to-end differentiability. When a PDE solver is part of a loss function, for instance in design optimization, parameter inference, or hybrid neural–physics model training, the underlying sparse linear solver must support automatic differentiation (AD) so that gradients can propagate backward through the linear solve to update design parameters, operator coefficients, or network weights. This places new demands on sparse linear solver software, which now needs to be not only fast and robust but also compatible with modern AD workflows.

The need to differentiate through PDE solvers is not new in spirit. PDE-constrained optimization is a mature field that combines numerical simulation, sensitivity analysis, and optimization to solve design, control, and inverse problems governed by PDEs. Classical approaches have relied on adjoint-based derivative computations [5]. In fluid dynamics, for example, adjoint design methods grew out of optimal-control ideas, including early work on optimum design by Pironneau [6] and pioneering adjoint-based design optimization by Jameson and collaborators [7–9]. These methods were later extended and adopted more widely in aerodynamic optimization [10–12]. Open-source frameworks such as SU2 [13] for computational fluid dynamics (CFD) and dolfin-adjoint [14] for finite-element PDE workflows have since made adjoint-based PDE-constrained optimization more practical in scientific computing.

Beyond these PDE-specific adjoint frameworks, general-purpose differentiable programming offers a more versatile route in modern scientific computing: building end-to-end differentiable workflows by composing discretized solvers, learnable neural networks, data assimilation, and uncertainty quantification within a single programming system. This hybrid neural–physics paradigm enables inverse problems, parameter inference, and optimal control through physics-as-architecture designs where the governing equations are embedded directly into the model architecture rather than enforced through loss functions [15–17]. JAX has emerged as a prominent ecosystem in this direction, combining array programming, AD, and just-in-time (JIT) compilation [18]. However, the bottleneck for such hybrid architectures is not the neural network, but the robust solution of large-scale sparse linear systems arising from the discretized PDEs within the backward pass. A prominent example arises in CFD when solving the pressure Poisson equation for the incompressible Navier–Stokes equations on complex computational meshes. The resulting linear systems are frequently ill-conditioned, and standard iterative solvers such as the conjugate gradient (CG) method or biconjugate gradient stabilized (BiCGSTAB) method often converge slowly or fail entirely on such systems. Algebraic multigrid (AMG) is one of the most effective approaches for these problems, using a hierarchy of progressively coarser representations to efficiently resolve errors across multiple scales [19]. Mature, high-performance AMG implementations are widely deployed across large-scale engineering and scientific computing, including hypre’s BoomerAMG [20,21], the PETSc toolkit with its native GAMG solver [22], the Trilinos ML/MueLu multigrid packages [23,24], and the header-only AMGCL library [25], all of which support MPI and, in several cases, GPU execution. These libraries are implemented in compiled languages (C/C++/Fortran) and are not integrated with AD or the JAX transformation system, and thus cannot be embedded directly into the end-to-end differentiable, GPU-native pipelines targeted here. Consequently, no high-performance AMG solver in the JAX ecosystem integrated both automatic differentiation and GPU-accelerated distributed execution.

The current Python software landscape addresses isolated parts of this problem but fails to provide a unified solution. For instance, PyAMG [26] offers a mature AMG implementation but operates exclusively on CPUs. The pyamgx package [27] provides a Pythonic interface to the GPU-accelerated AmgX backend API; however, it is restricted to single-GPU execution and lacks both AD and tracing capabilities. Similarly, petsc4py [28], which wraps the PETSc library [22] and supports GPU execution, AMG, and MPI, also lacks both AD and JIT compatibility, preventing its use in end-to-end differentiable pipelines. Conversely, differentiable sparse linear solvers within the JAX ecosystem, whether provided directly by JAX or by libraries like Lineax [29], currently lack support for AMG, leaving their iterative methods without the robust preconditioning required for ill-conditioned systems. A contemporaneous effort, AMJax [30], addresses this missing AMG capability from a complementary PyAMG/JAX direction: it builds the multigrid hierarchy on CPU with PyAMG and performs AMG solves in JAX. This design is lightweight and portable across JAX-supported backends after setup, but it does not target high-performance GPU-native AMG or distributed multi-GPU execution (a detailed performance comparison between JAX-AMG and AMJax is provided in Section 3.5). Similar limitations also arise beyond the JAX ecosystem: GPU-native differentiable frameworks such as NVIDIA Warp [31] likewise provide JIT compilation and AD, but expose only Krylov solvers with basic preconditioning rather than AMG, and are oriented toward implementing custom differentiable simulation kernels rather than serving as drop-in solvers within JAX-based pipelines. Together, these limitations motivate JAX-AMG as a bridge between high-performance GPU sparse solvers and differentiable Python workflows. It is built upon NVIDIA AmgX [32], a state-of-the-art library for GPU-accelerated sparse linear algebra that provides not only robust AMG algorithms but also a comprehensive suite of Krylov iterative methods. This integration brings scalable, differentiable sparse linear solvers and preconditioners to the JAX ecosystem.

Beyond simply exposing AmgX functionality in Python, JAX-AMG provides a unified interface deeply integrated with standard JAX workflows. It natively accepts JAX arrays and executes directly on GPU hardware, eliminating costly host-to-device memory transfers. Crucially, the library implements reverse-mode AD via discrete adjoint methods, seamlessly connecting the external sparse solver to JAX’s transformation system through primitives such as jax.grad. It also fully supports JIT compilation via jax.jit for tracing and optimizing end-to-end simulation pipelines. To address the scaling demands of massive computational grids, the package enables distributed multi-GPU execution using GPU-aware MPI. Furthermore, JAX-AMG introduces intelligent caching utilities that store and reuse solver setup phases, such as the construction of the multigrid hierarchy. This significantly reduces computational overhead during the repeated solves characteristic of time-stepping, inverse problems, and PDE-constrained optimization.

Table 1 summarizes how JAX-AMG addresses the gaps left by existing tools in the Python ecosystem. By unifying GPU acceleration, AMG, AD, JIT compilation, and MPI distribution in a single package, JAX-AMG resolves a critical bottleneck in modern differentiable scientific computing.

Table 1.

Feature comparison of Python sparse solver packages relevant to JAX-based workflows.

	GPU	AMG	AD	JIT	MPI
jax.scipy [18] (JAX-native)	✓	–	✓	✓	–
PyAMG [26]	–	✓	–	–	–
pyamgx [27] (via AmgX)	✓	✓	–	–	–
petsc4py [28] (via PETSc)	✓	✓	–	–	✓
Lineax [29]	✓	–	✓	✓	–
AMJax [30] (via PyAMG)	✓†	✓	✓	✓	–
JAX-AMG	✓	✓	✓	✓	✓

^†AMG hierarchy setup is CPU-side.

2. Software description

JAX-AMG provides GPU-accelerated sparse linear solvers with full support for AD and JIT compilation in JAX. This section describes the software architecture and its principal functionalities.

2.1. Software architecture

JAX-AMG is organized into three layers: a Python API layer, a JAX integration layer, and a native C++/CUDA backend, as illustrated in Fig. 1.

Fig. 1.

Software architecture of JAX-AMG.

Python API.

The main public interface is a single jaxamg.solve entry point, modeled after the native sparse solvers in the jax.scipy module. It accepts a matrix A and right-hand side b, along with an optional configuration dictionary for specifying solvers and preconditioners. Various input matrix formats are accepted, including JAX and SciPy sparse matrices, dense arrays, and callable linear operators. All inputs are normalized internally to the compressed sparse row (CSR) format. Solver configuration is specified via Python dictionaries, with the default being a BiCGSTAB solver with an AMG preconditioner.

JAX integration.

At the execution boundary, JAX-AMG registers CUDA foreign-function interface (FFI) targets with JAX’s XLA compiler and dispatches solves through jax.ffi.ffi_call, exposing the solver as a native JAX primitive that is fully compatible with jax.jit compilation. For AD, the library defines custom vector-Jacobian product (VJP) rules via jax.custom_vjp. Given a forward solve Ax = b and a scalar loss function $ℒ$, the backward pass computes the adjoint variable λ via the implicit function theorem,

$$A^{\text{T}}\lambda =\frac{\partial ℒ}{\partial x},$$ (1)

from which the gradients follow as

$$\frac{\partial ℒ}{\partial A}=-\lambda x^{\text{T}},\frac{\partial ℒ}{\partial b}=\lambda .$$ (2)

For symmetric matrices, the transpose solve is skipped. For nonsymmetric systems, the backward pass forms A^T using cuSPARSE on a single GPU or an all-to-all exchange in MPI mode, and then solves the adjoint system through the same AmgX backend with the same solver and preconditioner configurations.

Matrix-free operator support.

JAX-AMG accepts matrix-free linear operators, specified as callables x ↦ Ax, in addition to explicit sparse matrices. Because the AmgX backend requires explicit CSR data, the library materializes such an operator in two stages: it first recovers the operator’s sparsity pattern, then assembles the nonzero values. Value assembly uses graph-colored probing: a coloring of the sparsity graph groups structurally independent columns, so that all nonzero values are recovered with a number of operator evaluations equal to the number of colors rather than the matrix dimension [33,34]. This design targets sparse operators: the assembled CSR matrix must fit in memory, so operators that materialize as effectively large dense matrices may not be well suited to this approach.

The sparsity pattern itself is recovered by tracing: JAX-AMG interprets the operator’s computation graph and propagates a sparse connectivity structure through each primitive, recovering the sparsity pattern in a single pass with no operator evaluations [35]. Following the similar strategy as SparseConnectivityTracer.jl in the Julia ecosystem, this approach scales to large operators with negligible memory. When an operator cannot be traced—for instance, one with opaque custom-call primitives or data-dependent indexing—the library falls back to probing with one-hot basis vectors. This fallback is correct for any operator but materializes dense probe batches on the GPU and incurs $𝒪(N)$ evaluations.

When called outside JIT, this materialization happens automatically and is cached on the operator object. When a matrix-free operator appears inside JIT-compiled code, for example in an optimization loop with a parameterized operator, the sparsity pattern and coloring are precomputed once via cache_coloring and attached with with_cache, because they fix the matrix’s nonzero structure and color count, which JAX requires to be known at compile time. An illustrative example of this workflow is provided in Section 3.3.

C++/CUDA backend.

The C++ layer manages AmgX resources (matrices, vectors, solvers, configurations) and exposes four FFI handlers covering single- and multi-GPU modes in both single and double precision. An LRU cache retains AmgX solver handles keyed by a hash of the full sparsity pattern, together with the matrix sparsity dimensions, precision, and solver configuration. On a cache hit, the matrix coefficients are updated in place and the solver setup is repeated against the new values, while the cached solver resources and matrix structure from the initial solve are reused. This is substantially cheaper than a cold setup and remains correct whenever the sparsity pattern is unchanged, as in typical PDE-constrained optimization. If the pattern itself changes (e.g., from topology-changing mesh adaptation), the cache key changes and a fresh setup is triggered automatically.

Distributed execution.

In MPI mode, each rank holds a local matrix partition and right-hand side. The same jaxamg.solve interface is used, with additional arguments for the MPI communicator and global problem size. When this communicator spans multiple nodes, the same setup scales across them without modification. JAX-AMG does not currently support automatic sharding of a global JAX sparse matrix/vector through pjit or NamedSharding; users instead supply the local partitions directly, aided by helper utilities such as partition_csr_matrix and partition_operator. Communication is handled through mpi4jax [36], which provides MPI collectives as JAX-traceable primitives, with support for GPU-aware MPI when available. In the backward pass, for nonsymmetric A , the distributed transpose is assembled by an all-to-all exchange through mpi4jax that routes each local nonzero A_ij to the rank owning global row j (its position in A^T). AmgX then solves the resulting distributed A^T system. For matrix-value gradients in the distributed VJP, JAX-AMG exchanges only the remote solution entries referenced by each rank’s local sparse rows, avoiding an all-gather of the full global solution vector.

2.2. Software functionalities

The principal functionalities of JAX-AMG are:

• GPU-accelerated sparse solves using the full AmgX solver suite, including algebraic multigrid, Krylov subspace methods, and a variety of smoothers and preconditioners, all fully configurable via Python dictionaries compatible with AmgX’s configuration format.

• Reverse-mode automatic differentiation through linear solves via adjoint methods (Eq. 1), compatible with jax.grad. For symmetric matrices, the transpose solve is automatically skipped.

• JIT compilation via jax.jit for traced, compiled execution of end-to-end pipelines.

• Sequential jax.vmap compatibility for applying the solver over a batch of right-hand sides. Because the matrix is shared across the batch, solver caching amortizes the setup cost over the batch instead of repeating it for each right-hand side.

• Flexible input formats, including JAX and SciPy sparse matrices, dense arrays, and matrix-free operators.

• Automatic sparsity detection for matrix-free operators, recovering the sparsity pattern by tracing the operator’s computation graph in a single pass, with basis-vector probing as a fallback.

• Distributed multi-GPU execution via MPI, with support for GPU-aware MPI when available.

• Preconditioner export via jaxamg.make_preconditioner, returning a callable compatible with native JAX sparse linear solvers in jax.scipy and third-party libraries such as Lineax.

• Solver caching at both the Python and C++ levels to amortize setup costs across repeated solves with fixed sparsity patterns.

• Single and double precision arithmetic, with automatic promotion of mixed-precision inputs.

Solver behavior is controlled through a Python dictionary interface. Users need only specify the parameters they wish to change; all other settings inherit sensible defaults. The solver returns diagnostic information including iteration count, final residual, and convergence status. Optional file-based output provides detailed AMG hierarchy statistics and per-iteration residual histories.

3. Illustrative examples

We present four examples that demonstrate the core capabilities of JAX-AMG, followed by performance comparisons against native JAX solvers and other AMG libraries. Additional examples are available in the repository’s demo directory.

3.1. GPU-accelerated sparse solve with AMG preconditioning

The following example solves a 2D Poisson system on a 32 × 32 grid using BiCGSTAB with an AMG preconditioner:

1 import jaxamg

2 from jaxamg.matrices import poisson_matrix , rhs_ones

3

4 n = 32

5 A = poisson_matrix(n)

6 b = rhs_ones(n*n)

7

8 x, info = jaxamg.solve(A, b, config={

9    “solver”: “PBICGSTAB”,

10    “preconditioner”: {“solver”: “AMG”},

11    “tolerance”: le-6,

12 })

This example converges in 6 iterations, whereas the same solver without preconditioning requires 37 iterations, illustrating the effectiveness of AMG preconditioning for this class of problems. For more severely ill-conditioned systems, convergence without preconditioning may not be achievable at all.

3.2. Gradient-based optimization

A key use case for JAX-AMG is embedding linear solves inside gradient-based optimization loops. The following example recovers the unknown diagonal value d of a Toeplitz tridiagonal system such that the solution vector x matches a target solution x_target:

1 import jax

2 import jax.numpy as jnp

3 import jaxamg

4 from jaxamg.matrices import tridiagonal_matrix, rhs_ones

5

6 n = 32

7 true_diag = 4.0

8 init_diag = 10.0

9 lr = 0.1

10 max_iters = 200

11 tol = le-3

12

13 A_true = tridiagonal_matrix(n, diagonal_value=true_diag) 

14 b = rhs_ones (n)

15 x_target, _ = jaxamg.solve(A_true, b)

16

17 def loss(diag):

18    A = tridiagonal_matrix(n, diagonal_value=diag)

19    x, _ = jaxamg . solve (A , b)

20    return jnp.sum((x - x_target) ** 2)

21

22 grad_fn = jax.jit(jax.value_and_grad(loss))

23 diag = init_diag

24 for step in range (max_iters ) :

25    val, grad = grad_fn(diag)

26    diag -= Ir * grad

27    if jnp.linalg.norm ( grad ) < tol:

28       break

The gradient computation is JIT-compiled via jax.jit, with gradients obtained through jax.value_and_grad. In the backward pass, JAX-AMG automatically solves the adjoint system and computes gradients with respect to the matrix entries. Starting from an initial guess of d = 10.0, the optimizer recovers the true diagonal value d = 4.0 in fewer than 100 gradient steps under the specified learning rate and tolerance.

3.3. Gradient-based optimization with matrix-free operators

This example solves the same inverse problem as Section 3.2, recovering the diagonal value d of a Toeplitz tridiagonal system, but defines the operator A as a matrix-free callable rather than an explicit sparse matrix. Because the operator depends on d and changes at each optimization step, the sparsity pattern and graph coloring are precomputed via cache_coloring before the optimization loop and attached via with_cache inside the loss function, enabling JIT compilation. For this operator the sparsity pattern is recovered by tracing in a single pass with no operator evaluations.

1 import jax

2 import jax.numpy as jnp

3 import jaxamg

4 from jaxamg.matriceS import rhs_ones

5 

6 def A_operator(diag):

7    def matvec(x) :

8       y = -jnp.roll(x, 1) + diag * x - jnp.roll(x, −1)

9       y = y . at [0] . set (diag * x [0] - x [1] )

10      y = y . at [−1] . set (-x [−2] + diag * x[−l])

11      return y

12    return matvec

13 

14 n=32

15 true_diag = 4.0

16 init_diag = 10.0 

17 lr = 0.1

18 is max_iters = 200

19 tol = le-3

20

21 A_true = A_operator(true_diag)

22 b = rhs_ones (n)

23 x_target, _ = j axamg . solve (A_true , b)

24

25 diag = init_diag

26 coloring_cache = jaxamg.cache_coloring(A_operator(init_diag), shape=n)

27

28 def loss(diag):

29    A = jaxamg.with_cache(A_operator(diag), coloring = coloring_cache)

30    x, _ = j axamg . solve (A , b)

31    return jnp.sum((x - x.target) ** 2)

32

33 grad_fn = jax.jit(jax.value_and_grad(loss))

34

35 for step in range (max_iters ) :

36    val, grad = grad_fn(diag)

37    diag -= Ir * grad

38    if jnp . linalg . norm (grad) < tol :

39       break

The matrix-free formulation produces identical results to the explicit matrix case in Section 3.2, recovering d = 4.0 from an initial guess of d = 10.0 within the same number of gradient steps.

3.4. Distributed multi-GPU solve

JAX-AMG supports MPI-distributed solves with minimal changes to the single-GPU interface. The following example partitions a 2D Poisson problem across MPI ranks and solves it in parallel:

1 from mpi4py import MPI

2 import jaxamg

3 from jaxamg.matrices import poisson_matrix_distributed , rhs_linear

4 from jaxamg.mpi_utils import partition_vector , gather_vector

5

6 comm = MPI.C0MM_W0RLD

7 rank, nranks = comm.Get_rank () , comm.Get_size ()

8

9 n = 32

10 A_local , row_start , row_end = poisson_matrix_distributed(

11    n, n, rank, nranks)

12 b_local, _, _ = partition_vector(rhs_linear(n*n), rank, nranks)

13 x_local , info = jaxamg . solve (

14    A_local , b_local, comm=comm, nglobal=n*n,

15    partition_info=(row_start, row.end),

16    config={

17       “solver”: “PBICGSTAB”,

18       “preconditioner”: {“solver”: “MULTICOLOR.DILU”},

19       “communicator”: “MPI.DIRECT”})

20 x_global = gather_vector(x_local, comm, root=0)

21 jaxamg.finalize()

3.5. Performance comparison

Comparison with native JAX solvers.

To evaluate the performance of JAX-AMG, we first compare it against the native JAX CG and BiCGSTAB sparse linear solvers on a tridiagonal system of dimension n = 10⁷ . To isolate the effect of the solver backend from that of preconditioning or other algorithmic differences, we deliberately choose a well-conditioned system that the native JAX solvers can handle without preconditioning, and configure JAX-AMG to use the same unpreconditioned CG and BiCGSTAB methods. Fig. 2(a) compares wall-clock time for a single forward solve, Fig. 2(b) reports JIT compilation time for the gradient computation, and Fig. 2(c) shows the total optimization time over 20 iterations. In all three cases, JAX-AMG outperforms native JAX solvers, demonstrating the performance advantage of the GPU-optimized AmgX backend.

Fig. 2.

Performance comparison between JAX-AMG and native JAX solvers (CG and BiCGSTAB) on a tridiagonal system with n = 10⁷ unknowns. (a) Forward solve wall-clock time. (b) JIT compilation time for gradient calculation. (c) Total optimization time for 20 iterations. All solvers use the same unpreconditioned configuration. Benchmarked on an NVIDIA L40 GPU.

Comparison with established AMG libraries.

We next benchmark JAX-AMG against the native AmgX backend and two widely used AMG libraries (BoomerAMG and AMGCL) on three large, structurally distinct sparse systems from the SuiteSparse Matrix Collection [37]: G3_circuit (circuit simulation; structurally and value-symmetric), tmt_unsym (electromagnetic simulation; structurally symmetric but value-nonsymmetric), and torso3 (electrophysiological modeling; structurally nonsymmetric). These systems span different application domains and sparsity patterns. Each system is solved with a Krylov iterative method preconditioned by classical Ruge–Stüben and aggregation-based AMG on 1, 2, and 4 GPUs using the same configurations. We report the cost of a warm resetup-and-solve, because this is the operation typically repeated within optimization and implicit time-stepping loops, where the matrix entries change but the sparsity pattern remains fixed, and therefore dominates the end-to-end cost. Fig. 3 presents the six resulting configurations (three matrices × classical/aggregation AMG), comparing native AmgX and JAX-AMG across the three GPU counts, together with BoomerAMG [21] for the classical AMG panels and AMGCL [25] for the aggregation AMG panels (BoomerAMG is classical-only, whereas AMGCL’s MPI path is aggregation-only). Because AMG libraries expose different configurable options, we matched the solver settings as closely as possible when exact equivalence was not available. These benchmarks primarily compare solver implementations, and the GPU-count trends should be interpreted with care because distributed resetup and communication overheads can produce non-monotonic timings for these matrix sizes. Across all cases, JAX-AMG closely tracks native AmgX, confirming that the JAX FFI introduces negligible overhead on the repeated solve path. On multiple GPUs, JAX-AMG is sometimes even slightly faster than the native backend. Since both backends issue identical AmgX operations, this is not an algorithmic advantage but a runtime effect, likely related to how the JAX/XLA runtime schedules host work around MPI progress and communication overlap during distributed resetup. Relative to other established libraries, AmgX/JAX-AMG is competitive, although no AMG library is consistently superior. We note that AmgX performs resetup entirely on the GPU, whereas BoomerAMG offers no lightweight resetup and must rebuild its hierarchy from scratch, and AMGCL performs its rebuild on the host CPU.

Fig. 3.

Performance comparison of classical and aggregation AMG solve times for three sparse matrices across 1, 2, and 4 GPUs. Panels (a–c) show classical AMG results (JAX-AMG, AmgX, and BoomerAMG) for matrices G3_circuit, tmt_unsym, and torso3, respectively. Panels (d–f) show aggregation AMG results (JAX-AMG, AmgX, and AMGCL) for the same three matrices. Within each panel, the x-axis represents GPU count (1, 2, or 4 GPUs), and bars show performance for each solver. Reported timings measure a warm resetup-and-solve step on NVIDIA L40 GPUs.

Comparison with AMJax.

We next compare JAX-AMG with AMJax [30], a JAX-oriented AMG solver built on PyAMG. In this comparison, AMJax is used as a preconditioner for JAX-native Krylov solvers, with the solver and preconditioner configurations matched as closely as the two libraries allow. The comparison uses the same three sparse matrices considered above. Since AMJax does not provide an AmgX-style resetup operation, we compare fixed-hierarchy repeated solves and run JAX- AMG with reuse_setup = True to reuse the existing hierarchy without resetup. The experiments are limited to a single GPU because AMJax does not support distributed execution. The results are shown in Fig. 4. On G3_circuit and tmt_unsym, AMJax converges but is substantially slower than JAX-AMG for both forward and gradient computations. In the most challenging case, torso3, JAX-AMG still converges, whereas AMJax diverges under the closest matched configuration, and additional tuning attempts also failed to produce valid forward or gradient solves. While AMJax offers a more portable JAX-based workflow, JAX-AMG achieves stronger robustness and performance for the large-scale GPU AMG workloads considered here, primarily due to its efficient AmgX backend and richer solver configurations.

Fig. 4.

Performance comparison of JAX-AMG and AMJax on three benchmark matrices (G3_circuit, tmt_unsym, and torso3). Each panel reports wall-clock times for the first forward solve, subsequent forward solve, first gradient solve, and subsequent gradient solve. Hatched bars overlaid with a red cross marker denote AMJax runs that failed to converge. Benchmarked on an NVIDIA L40 GPU with 48 GB of memory. (For interpretation of the references to color in this figure legend, the reader is referred to the web version of this article.)

Comparison with PETSc for CFD application.

Finally, we test JAX-AMG on a practical problem where preconditioning is essential, integrating it as the pressure Poisson solver in Diff-FlowFSI [38], our in-house differentiable incompressible CFD solver, to simulate turbulent channel flow at friction Reynolds number Re_T = 390, where unpreconditioned BiCGSTAB fails to converge. The computational domain is L_x × L_y × L_z = 2 π × 1 × π , discretized on a 100 × 260 × 256 grid with wall-normal stretching. We compare JAX-AMG against PETSc [22,28] both configured with BiCGSTAB and multigrid preconditioning. In this stretched-grid finite-volume case, the stored pressure Poisson matrix is nonsymmetric, so BiCGSTAB is used. Compatibility and zero-mean constraints are enforced for the periodic/Neumann pressure setup to remove the constant null mode. As shown in Table 2, JAX-AMG requires fewer iterations and lower per-GPU memory than PETSc under the same convergence tolerance, at the cost of a moderate increase in solve time. On one and two GPUs, JAX-AMG converges in 5–6 iterations compared with 15 for PETSc, while reducing per-GPU memory by about 29% and 41%, respectively. The two-GPU timings correspond to fixed-problem-size distributed runs, where communication and distributed-solver overheads can offset the reduced local workload. Fig. 5 confirms that the resulting turbulence statistics are in close agreement. Small discrepancies in second-order statistics between the two results are expected because the solvers use different AMG implementations and configuration options. Crucially, unlike PETSc, JAX-AMG retains full compatibility with JAX’s AD and JIT compilation, enabling gradient-based optimization through the same solver without breaking the differentiable pipeline.

Table 2.

Computational cost per Poisson solve comparing JAX-AMG and PETSc in the Diff-FlowFSI solver for incompressible turbulent channel flow at Re_τ = 390. Benchmarked on NVIDIA RTX A6000 GPUs.

Item	Single GPU		2 GPUs
	PETSc	JAX-AMG	PETSc	JAX-AMG
Time (s)	2.90	3.88	2.80	4.25
Memory (MiB)/GPU	11,776	8348	8758	5130
Convergence iteration	15	5	15	6
Convergence error	10−5	10−5	10−5	10−5

Fig. 5.

Velocity statistics from Diff-FlowFSI using JAX-AMG and PETSc as the Poisson solver for turbulent channel flow at Re_τ = 390: (a) mean streamwise velocity; (b–d) root-mean-square of streamwise, wall-normal, and spanwise velocities, respectively.

4. Impact

JAX-AMG fills a specific gap in the scientific computing toolchain: it is, to our knowledge, the first package to combine GPU-accelerated AMG with AD, JIT compilation, and distributed execution within the JAX ecosystem. A growing number of JAX-based differentiable simulation frameworks have been developed across various domains [39], including finite element methods [40], CFD [38,41], and nanoscale heat transfer [42], among others. Many of these frameworks rely on sparse linear solves as a core computational kernel, yet the solver options currently available in JAX are limited: native solvers lack robust preconditioning, while external GPU solver libraries lack differentiability and JIT compatibility. JAX-AMG resolves this tension, and its impact extends across several areas.

The most direct impact is on PDE-constrained optimization and inverse problems. These workflows require differentiating through repeated sparse solves, often involving ill-conditioned systems where unpreconditioned Krylov methods either converge slowly or fail entirely. Because JAX-AMG provides AMG-preconditioned solves with gradient support, problems such as design optimization, optimal control, and spatially varying parameter inference can now be tackled end-to-end within JAX. Prior to JAX-AMG, researchers in the JAX ecosystem faced a choice between native solvers that lacked robust preconditioning and external solvers that broke the differentiable pipeline. JAX-AMG eliminates this trade-off.

A second area of impact is large-scale differentiable simulation. The MPI-distributed backend partitions problems across multiple GPUs while retaining full gradient support, extending differentiable optimization to problem scales that exceed single-GPU memory. This is particularly relevant for high-resolution 3D PDE systems—such as turbulent flow or full-domain cardiac hemodynamics simulations—where single-GPU memory constraints would otherwise preclude gradient-based inversion or optimal control.

Beyond enabling new research, JAX-AMG improves the efficiency and flexibility of existing JAX-based workflows. Even for well-conditioned systems where preconditioning is unnecessary, the GPU-optimized AmgX backend achieves considerable speedups over native JAX solvers as demonstrated in Section 3.5. The jaxamg.make_preconditioner function allows JAX-AMG to serve as a preconditioning backend for native JAX or other third-party iterative solvers, bringing AMG preconditioning to solver pipelines that previously had access only to simple preconditioners. The unified Python interface, covering solver configuration, gradient computation, JIT compilation, and MPI distribution, reduces engineering effort and improves reproducibility compared to workflows that stitch together tools across multiple languages and runtimes.

JAX-AMG is publicly available at https://github.com/jx-wang-s-group/JAX-AMG under the Apache 2.0 license. It has already been integrated into Diff-FlowFSI [38], a differentiable CFD platform for turbulent flow and fluid-structure interaction, where it serves as the pressure Poisson solver enabling gradient-based optimization through high-Reynolds-number simulations. It has similarly been adopted in JAX-BTE [42], a differentiable solver for nanoscale phonon transport, demonstrating JAX-AMG’s applicability across diverse PDE domains. These integrations validate JAX-AMG as a practical computational substrate for differentiable simulation in both fluid mechanics and thermal transport.

5. Conclusions

JAX-AMG bridges the gap between high-performance GPU sparse solvers and the JAX differentiable programming ecosystem. By wrapping the full NVIDIA AmgX solver suite, including AMG and Krylov methods, as a native JAX primitive, it brings robust GPU-accelerated sparse linear algebra into differentiable workflows with support for reverse-mode AD, JIT compilation, jax.vmap compatibility, and MPI-distributed execution. These capabilities make JAX-AMG a practical foundation for the growing ecosystem of JAX-based differentiable simulation frameworks that rely on sparse linear solves as a core computational kernel.