JAXMg

JAXMg brings distributed matrix solvers to JAX, allowing calculations to run across multiple GPUs and nodes. This enables solvers to scale to matrices near the combined memory capacity of the available GPUs, far beyond native JAX routines, while retaining a familiar JAX interface.

JAXMg currently provides a jittable API for the following routines:

• potrs: Solves the system of linear equations \(Ax=B\), where \(A\) is an \(N \times N\) symmetric or Hermitian positive-definite matrix, using a Cholesky decomposition. It can also return the log determinant of \(A\).
• lu_solve: Solves the system of linear equations \(Ax=B\), where \(A\) is an \(N \times N\) general nonsingular matrix, using a pivoted LU decomposition.
• syevd: Computes the eigenvalues and optional eigenvectors of an \(N \times N\) symmetric or Hermitian matrix.

How JAXMg works

JAXMg connects JAX to NVIDIA's distributed cuSOLVERMp routines through a native C++/CUDA backend.

Supply a JAX matrix sharded over a two-dimensional device mesh, and JAXMg handles the local memory-layout conversion, redistribution into cuSOLVERMp's 2D block-cyclic layout, distributed numerical computation, and restoration of the result to its original JAX layout. The matrix data remains GPU-resident throughout, with in-place transformations and bounded scratch storage minimizing memory overhead.

The operations are implemented using:

• potrs: cusolverMpPotrf and cusolverMpPotrs
• lu_solve: cusolverMpGetrf and cusolverMpGetrs
• syevd: cusolverMpSyevd

For more details, see the API reference and the accompanying paper.

Deprecation notice: JAXMg 0.x, based on NVIDIA's deprecated cuSOLVERMg backend, has been superseded by JAXMg 1.0. Version 1.0 requires one Python process per GPU and removes potri; the previous API remains available with pip install "jaxmg<1" and in the 0.0.9 release.