Related Experiment Video
Updated: Oct 26, 2026

Swin-PSAxialNet: An Efficient Multi-Organ Segmentation Technique
Published on: July 5, 2024
JAX-AMG: A GPU-accelerated differentiable sparse linear solver library for JAX
Yi Liu1,2, Xiantao Fan1, Jian-Xun Wang1,2
1Sibley School of Mechanical and Aerospace Engineering, Cornell University, Ithaca, NY, USA.
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.
Related Concept Videos
Lagrange Multipliers: Problem Solving
Linear Approximations
Gauss's Law: Problem-Solving
Ampere-Maxwell's Law: Problem-Solving
To solve the problem, we can use the equations from the analysis of an RC circuit and Maxwell's version of Ampère's law.
For the first part of the problem,...
Statically Indeterminate Problem Solving
Linear Differential Equations