Year: 2026
Type: preprint
Abstract: Solving large dense linear systems and eigenvalue problems is a core requirement in many areas of scientific computing, but scaling these operations beyond a single GPU remains challenging within modern programming frameworks. While highly optimized multi-GPU solver libraries exist, they are typically difficult to integrate into composable, just-in-time (JIT) compiled Python workflows. JAXMg provides distributed dense linear algebra for JAX, enabling linear solves and decompositions for matrices that exceed single-GPU memory limits. By interfacing JAX with NVIDIA's cuSOLVERMp through an XLA Foreign Function Interface, JAXMg exposes distributed GPU routines as JIT-compatible JAX primitives. This design allows scalable linear algebra to be embedded directly within JAX programs, preserving composability with JAX transformations and enabling multi-GPU and multi-node execution in end-to-end scientific workflows.
Source: arXiv (Cornell University)
Authors: Tutt, Jacob, Wiersema, Roeland
Institutions: University of Cambridge, Flatiron Institute
Subfield: Hardware and Architecture
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
Awards: G2912470688, G963556061
OpenAlex ID: W7125418037
Cites:
Cited by:
Related to: