Examples

We provide a complete walkthrough of the JAXMg workflow, from launching a distributed JAX program to running a cuSOLVERMp routine.

Ready-to-run scripts are provided in the repository's examples directory. They establish the mesh and input sharding once, after which each solver infers that layout directly from its input matrix.

1. Configure Distributed Execution

Start by initializing distributed JAX, assigning one GPU to each process, and constructing the device mesh used by the matrix.

Configure distributed execution

2. Choose a tile size

Choose a cuSOLVERMp tile size \(T_A\) that divides both dimensions of each local matrix shard. Choosing an appropriate tile size can substantially improve runtime and reduce memory use.

Choose a tile size

3. Run a solver

• Cholesky solve solves a symmetric or Hermitian positive-definite linear system with potrs.
• LU solve solves a general nonsingular linear system with lu_solve.
• Least-squares solve solves an overdetermined rectangular system with least_squares.
• Reduced QR decomposition computes an orthonormal basis and upper-triangular factor with qr.
• Symmetric or Hermitian eigensolve computes eigenvalues and optional eigenvectors with syevd.
• Singular-value decomposition computes the singular values and optional left and right singular vectors of a real or complex matrix with gesvd.
• Polar decomposition computes the polar factor and optional Hermitian factor with polar.