jaxmg.potrs

potrs is the high-level solver for symmetric or Hermitian positive-definite matrices. It uses cuSOLVERMp Potrf followed by Potrs, with memory-layout conversion, redistribution, and solver execution in one fused native call.

Use potrs_jit_ctx when the solve is part of a larger caller-owned jax.jit computation. Inside an existing jax.shard_map, use potrs_shardmap_ctx with local matrix and RHS blocks instead. Both use the same distributed solver.

Context API migration

The previous global-array potrs_shardmap_ctx is now named potrs_jit_ctx, with the same arguments and outputs.

All three 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, donate=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 sharded over a one- or two-axis device mesh, for example with P(<row_axis>) or 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 or AbstractMesh	JAX mesh used for jax.shard_map. If omitted, read off the sharding of a (its type inside jax.jit), or taken from the context mesh.	None
matrix_specs	PartitionSpec or tuple / list[PartitionSpec]	PartitionSpec describing the matrix sharding. If omitted, read off the sharding of a, defaulting to the mesh axes in order.	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
donate	bool	If True (default) the input buffers may be donated to the native call for zero-copy execution, which means they are deleted and cannot be used again. Pass False to preserve them, at the cost of keeping the original and working buffers in memory simultaneously.	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

• Unless donate=False, the a and b buffers are donated for zero-copy interaction with the native library, so they are deleted and cannot be used after the call.
• 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 solution and logdet are filled with NaN and status is non-zero.

jaxmg.potrs_jit_ctx(a, b, T_A, mesh=None, matrix_specs=None, *, in_specs=None, return_logdet=False, pad=True, _prepared=None)

Solve A x = B on global arrays inside a caller-owned jax.jit.

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.

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 sharded over a one- or two-axis device mesh, for example with P(<row_axis>) or 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 or AbstractMesh	JAX mesh used for jax.shard_map. If omitted, read off the sharding of a (its type inside jax.jit), or taken from the context mesh.	None
matrix_specs	PartitionSpec or tuple / list[PartitionSpec]	PartitionSpec describing the matrix sharding. If omitted, read off the sharding of a, defaulting to the mesh axes in order.	None
in_specs	P | Tuple[P] | List[P] | None	Backwards-compatible alias for matrix_specs.	None
return_logdet	bool	If True return the replicated Cholesky log determinant between the solution and status outputs. 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:

Name	Type	Description
tuple	Union[Tuple[Array, Array, Array], Tuple[Array, Array, Array, Array]]	(a_work, x, status) or
	Union[Tuple[Array, Array, Array], Tuple[Array, Array, Array, Array]]	(a_work, x, logdet, status). a_work is the padded matrix work
	Union[Tuple[Array, Array, Array], Tuple[Array, Array, Array, Array]]	buffer required for donation, x retains the JAX-facing layout of
	Union[Tuple[Array, Array, Array], Tuple[Array, 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

• This function intentionally returns a_work. Public :func:potrs discards that buffer for convenience, but an outer jax.jit(..., donate_argnums=(0, 1)) can only donate a if an A-sized output is returned. Callers that want donation must keep a_work in the jitted function's returned pytree.
• 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.

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

Solve A x = B on local blocks inside a caller's jax.shard_map.

Use this interface when a distributed solve forms part of a larger shard-local calculation. The enclosing mesh must contain exactly the matrix axes. A follows matrix_specs; B must share its row sharding, but may be distributed or replicated over the process-column axis.

Parameters:

Name	Type	Description	Default
a	Array	Local block of a 2D, symmetric positive-definite matrix.	required
b	Array	Local block of the solve input, sharded like a or replicated over the column axis. A local vector is treated as a one-column block.	required
T_A	int	Square tile width used by cuSOLVERMp. Each local shard dimension must be a multiple of T_A after padding.	required
matrix_specs	PartitionSpec or tuple / list[PartitionSpec]	PartitionSpec of a in the enclosing shard_map. If omitted, defaults to the context mesh axes in order.	None
in_specs	P | Tuple[P] | List[P] | None	Backwards-compatible alias for matrix_specs.	None
return_logdet	bool	If True return the replicated Cholesky log determinant between the solution and status outputs. Default is False.	False
pad	bool	If True (default) pad each local block so its dimensions are compatible with T_A; if False the caller must ensure shapes already match the kernel's requirements.	True

Returns:

Name	Type	Description
tuple	Union[Tuple[Array, Array, Array], Tuple[Array, Array, Array, Array]]	(a_work, x, status) or (a_work, x, logdet, status).
	Union[Tuple[Array, Array, Array], Tuple[Array, Array, Array, Array]]	a_work is the padded local matrix work buffer required for
	Union[Tuple[Array, Array, Array], Tuple[Array, Array, Array, Array]]	donation, x is the local solution block, logdet is a real
	Union[Tuple[Array, Array, Array], Tuple[Array, Array, Array, Array]]	scalar that is identical on every rank, and status is this rank's
	Union[Tuple[Array, Array, Array], Tuple[Array, Array, Array, Array]]	native diagnostic vector.

Raises:

Type	Description
TypeError	If dtypes or PartitionSpec inputs are unsupported.
ValueError	If called outside jax.shard_map, or if shapes, tile sizes, or mesh layouts are incompatible.

Notes

• When B is replicated over process columns, its copies are solved as duplicate right-hand sides using one distributed factorization.
• Return a_work from the outer jitted function when donating A; it provides the corresponding matrix-sized output buffer.