jaxmg.potrs¤

potrs is the high-level Cholesky solve interface. It validates the arrays and mesh, applies tile-capacity padding when required, runs the internally compiled fused backend, and returns the solution in the JAX-facing layout.

Use potrs for a direct solve. Use potrs_shardmap_ctx only when the solve must be embedded inside a larger caller-owned jax.jit computation. Both interfaces run the same native solver pipeline; _shardmap_ctx means that the caller owns the surrounding compilation context, not that it selects a different solver or CUDA context.

Both interfaces accept return_logdet=True. Since the factorization produces \(A=LL^H\), the backend computes

\[ \log\det(A) = 2\sum_i \log |L_{ii}| \]

directly from the distributed Cholesky factor.

`jaxmg.potrs(a, b, T_A, mesh=None, matrix_specs=None, *, in_specs=None, return_status=False, return_logdet=False, pad=True)` ¤

Solve the linear system A x = B using the multi-GPU potrs native kernel.

This is the high-level JAXMg Cholesky-solve entry point. It prepares a block-sharded JAX array for cuSOLVERMp, calls the fused native backend, and returns the solution in the same JAX-facing layout as b.

Note

If a local shard dimension is not divisible by T_A, pad=True allocates additional tile-aligned capacity before the native call. Choosing a tile size that divides the local dimensions avoids this allocation. Performance depends on the matrix size, process grid, and tile size.

Parameters:

Name	Type	Description	Default
`a`	`Array`	2D, symmetric positive-definite input matrix. Expected to be sharded across a 2D mesh with a matrix `PartitionSpec` such as `P(<row_axis>, <col_axis>)`.	required
`b`	`Array`	1D or 2D solve input. A vector is treated as an `N x 1` matrix.	required
`T_A`	`int`	Square tile width used by cuSOLVERMp. Each local shard dimension must be a multiple of `T_A` after padding.	required
`mesh`	`Mesh`	JAX mesh used for `jax.shard_map`. If omitted, inferred from `a.sharding.mesh`.	`None`
`matrix_specs`	`PartitionSpec or tuple / list[PartitionSpec]`	PartitionSpec describing the matrix sharding. If omitted, inferred from `a.sharding.spec`.	`None`
`in_specs`	`P \| Tuple[P] \| List[P] \| None`	Backwards-compatible alias for `matrix_specs`.	`None`
`return_status`	`bool`	If True return `(x, status)` where `status` is the native per-rank diagnostic vector. If False return `x` only. Default is False.	`False`
`return_logdet`	`bool`	If True also return `log(det(A))` computed from the distributed Cholesky factor. Default is False.	`False`
`pad`	`bool`	If True (default) apply per-device padding so each local shard length is compatible with `T_A`; if False the caller must ensure shapes already match the kernel's requirements.	`True`

Returns:

Type	Description
`Union[Array, Tuple[Array, Array], Tuple[Array, Array, Array]]`	One of `x`, `(x, status)`, `(x, logdet)`, or
`Union[Array, Tuple[Array, Array], Tuple[Array, Array, Array]]`	`(x, logdet, status)`. The solution retains the JAX-facing layout of
`Union[Array, Tuple[Array, Array], Tuple[Array, Array, Array]]`	`b` and `logdet` is a replicated real scalar.

Raises:

Type	Description
`TypeError`	If dtypes or `PartitionSpec` inputs are unsupported.
`ValueError`	If shapes, tile sizes, or mesh layouts are incompatible.

Notes

The FFI call may donate the a and b buffers for zero-copy interaction with the native library.
Native code converts row-major JAX local storage to cuSOLVERMp's column-major local layout, redistributes to 2D block-cyclic layout, calls cusolverMpPotrf/cusolverMpPotrs, and redistributes the result back.
If the native solver fails the returned solution may contain NaNs and status will be non-zero.

`jaxmg.potrs_shardmap_ctx(a, b, T_A, mesh=None, matrix_specs=None, *, in_specs=None, return_logdet=False, pad=True)` ¤

Solve A x = B while exposing the donated matrix work buffer.

This helper is the lower-level variant of :func:jaxmg.potrs intended for contexts where the caller wants to control the outer jax.jit boundary. It performs the same validation, local padding, shard-map construction, and fused cuSOLVERMp FFI call as the public solver, but it does not wrap the pipeline in an internal jax.jit. Instead, it returns the native matrix work buffer alongside the solution so an outer JIT can donate a into an A-sized output.