jaxmg.potrf

potrf computes the lower Cholesky factor of a symmetric or Hermitian positive-definite matrix \(A\):

\[ A = L L^{\dagger}. \]

Use potrf for a direct factorization. Use potrf_jit_ctx when the factorization is part of a larger caller-owned jax.jit. The factor aliases the input work buffer, so the context interface returns \(L\) directly rather than a separate opaque work result.

Inside an existing jax.shard_map, use potrf_shardmap_ctx with local matrix blocks. It returns local factor blocks and per-rank status.

Context API migration

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

jaxmg.potrf(a, T_A, mesh=None, matrix_specs=None, *, in_specs=None, return_status=False, pad=True, donate=True)

Compute the distributed Cholesky factor of a positive-definite matrix.

For a symmetric or Hermitian positive-definite n x n matrix A, returns the lower-triangular matrix L such that A = L @ L^H. The strict upper triangle of L is zero, as in :func:jax.numpy.linalg.cholesky.

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
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 (L, status) where status is the native per-rank diagnostic vector. If False return L only. 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 buffer may be donated and overwritten by L, which means it is deleted and cannot be used again. Pass False to preserve it, at the cost of an additional A-sized allocation.	True

Returns:

Type	Description
Union[Array, Tuple[Array, Array]]	Array or (Array, Array): The Cholesky factor L, restored to replicated sharding when a is replicated and otherwise sharded like a. If return_status=True also return the native solver status.

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 buffer is donated for zero-copy interaction with the native library, so it is 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, and redistributes the factor back.
• If the native solver fails, for example because a is not positive definite, L is filled with NaN and status is non-zero.

---

jaxmg.potrf_jit_ctx(a, T_A, mesh=None, matrix_specs=None, *, in_specs=None, pad=True, _prepared=None)

Compute the Cholesky factor inside a caller-owned jax.jit.

This interface performs the same validation, redistribution, and native computation as :func:jaxmg.potrf, but leaves the outer compilation and donation boundary to the caller. The native factor output aliases the input work buffer, so an outer jax.jit(..., donate_argnums=(0,)) can donate a into L.

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
T_A	int	Square tile width used by cuSOLVERMp.	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
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
Array	(L, status). The status vector is always returned so an enclosing
Array	compiled function can propagate native diagnostics. L is restored
Tuple[Array, Array]	when a is replicated; otherwise it uses matrix_specs.

Raises:

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

Notes

• If the native solver fails, L is filled with NaN and the status is non-zero.

jaxmg.potrf_shardmap_ctx(a, T_A, matrix_specs=None, *, in_specs=None, pad=True)

Compute local Cholesky-factor blocks inside a caller's jax.shard_map.

The enclosing mesh contains the matrix axes. The same native factorization is used by :func:potrf and :func:potrf_jit_ctx; only the caller's compilation and sharding context differs.

Parameters:

Name	Type	Description	Default
a	Array	Local block of a symmetric or Hermitian positive-definite matrix.	required
T_A	int	Square tile width used by cuSOLVERMp.	required
matrix_specs	P | Tuple[P] | List[P] | None	A's partition specification in the enclosing shard map; defaults to the context mesh axes in order.	None
in_specs	P | Tuple[P] | List[P] | None	Backwards-compatible alias for matrix_specs.	None
pad	bool	Allow tile-capacity padding; otherwise require aligned blocks.	True

Returns:

Type	Description
Array	(L, status): the unpadded local factor block and this rank's
Array	diagnostic vector. The factor reuses the input work buffer, so no
Tuple[Array, Array]	separate a_work output is needed for donation.

Raises:

Type	Description
TypeError	If the dtype or partition specification is unsupported.
ValueError	If called outside shard_map or the global matrix is invalid.