jaxmg.getrf

getrf computes the pivoted LU factorization of a square nonsingular matrix \(A\):

\[ P A = L U. \]

The result lu stores the strictly lower part of the unit-lower factor \(L\) and the upper-triangular factor \(U\). The replicated vector piv records the zero-based row interchanges.

Use getrf for a direct factorization. Use getrf_jit_ctx when the factorization is part of a larger caller-owned jax.jit. The packed factor aliases the input work buffer, so no separate opaque work result is needed.

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

Compute the distributed pivoted LU factorization of a square matrix.

For a nonsingular n x n matrix A, computes P @ A = L @ U with partial pivoting. The outputs follow :func:jax.scipy.linalg.lu_factor: lu holds the strictly lower part of the unit-lower L and the upper triangle of U, and piv holds 0-based pivot indices, meaning row i was interchanged with row piv[i]. Use :func:jax.lax.linalg.lu_pivots_to_permutation to obtain the permutation.

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, nonsingular 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 (lu, piv, status) where status is the native per-rank diagnostic vector. If False return (lu, piv) 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 lu, 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[Tuple[Array, Array], Tuple[Array, Array, Array]]	(lu, piv) by default, or (lu, piv, status) when
Union[Tuple[Array, Array], Tuple[Array, Array, Array]]	return_status=True. lu is restored to replicated sharding when
Union[Tuple[Array, Array], Tuple[Array, Array, Array]]	a is replicated and otherwise sharded like a. piv is a
Union[Tuple[Array, Array], Tuple[Array, Array, Array]]	replicated int32 vector of length n.

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 cusolverMpGetrf, and redistributes the factor back.
• If the native solver fails, including when U has an exactly zero diagonal entry, lu is filled with NaN, piv is filled with -1, and status is non-zero.

---

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

Compute the pivoted LU factorization inside a caller-owned jax.jit.

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

Parameters:

Name	Type	Description	Default
a	Array	2D, nonsingular 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	(lu, piv, status). The status vector is always returned so an
Array	enclosing compiled function can propagate native diagnostics. lu
Array	is restored when a is replicated; otherwise it uses
Tuple[Array, Array, Array]	matrix_specs. piv is a replicated int32 vector.

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, lu is filled with NaN, piv is filled with -1, and the status is non-zero.

Use getrf_shardmap_ctx inside an existing jax.shard_map. It accepts a local matrix block and returns a local factor block, replicated pivots and per-rank status.

Context API migration

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

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

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

The enclosing mesh contains the matrix axes. Factor blocks are local, but the native backend returns the complete pivot vector on every rank.

Parameters:

Name	Type	Description	Default
a	Array	Local block of a square, nonsingular global 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	(lu, piv, status): the unpadded local factor block, replicated
Array	zero-based pivots and this rank's diagnostic vector. The factor reuses
Array	the input work buffer; no separate a_work output is needed.

Raises:

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