jaxmg.lu_solve

lu_solve is the high-level solver for general nonsingular matrices. It uses cuSOLVERMp Getrf followed by Getrs, including distributed pivot storage and the same fused memory-distribution path used by potrs.

Use lu_solve_shardmap_ctx when the solve is part of a larger caller-owned jax.jit computation.

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

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

This is the high-level JAXMg general linear-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, nonsingular 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
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]]	Array or (Array, Array): The solution x in the same JAX-facing block-sharded layout as b. 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

• 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 cusolverMpGetrf/cusolverMpGetrs, and redistributes the result back.
• If the native solver fails the returned solution may contain NaNs and status will be non-zero.

jaxmg.lu_solve_shardmap_ctx(a, b, T_A, mesh=None, matrix_specs=None, *, in_specs=None, pad=True)

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

This helper is the lower-level variant of :func:jaxmg.lu_solve 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.

Parameters:

Name	Type	Description	Default
a	Array	2D, nonsingular 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.	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
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	Array	(a_work, x, status) where a_work is the padded matrix
	Array	work buffer returned by the native FFI call, x is the solution in
	Array	the same JAX-facing layout as b, and status is the native
	Tuple[Array, Array, Array]	per-rank diagnostic vector.