jaxmg.least_squares

least_squares solves an overdetermined linear system using distributed QR factorization:

\[ X = \underset{X}{\operatorname{argmin}}\;\lVert A X-B\rVert_2, \qquad A\in\mathbb{F}^{M\times N},\quad M\geq N. \]

The solve input \(B\) may be a vector or a matrix containing multiple columns. The returned solution has shape (N,) or (N, K), respectively. Each process-grid column must own at least one block-cyclic tile of \(B\); a vector input therefore requires a process grid with one column.

Use least_squares for a direct solve. Use least_squares_jit_ctx when the solve is part of a larger caller-owned jax.jit. Since cuSOLVERMp overwrites both distributed inputs, the context interface returns a_work and b_work as shape-compatible alias targets for donated arguments.

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

Solve min_X ||A X - B||_2 with distributed QR factorization.

The current implementation supports overdetermined or square systems with M >= N. For A of shape (M, N) and B of shape (M, K), the returned solution has shape (N, K). A rank-1 B is accepted and produces a rank-1 solution. Every process-grid column must own at least one block-cyclic tile of B.

Parameters:

Name	Type	Description	Default
a	Array	Rank-2 input matrix sharded over a one- or two-axis device mesh.	required
b	Array	Rank-1 or rank-2 solve input with M rows.	required
T_A	int	Square cuSOLVERMp tile width.	required
mesh	Mesh or AbstractMesh	JAX mesh used by jax.shard_map. If omitted, read from the sharding of a (its type inside jax.jit), or taken from the context mesh.	None
matrix_specs	PartitionSpec or tuple / list[PartitionSpec]	Rank-2 matrix sharding. If omitted, read from 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).	False
pad	bool	If True (default), add tile-aligned local capacity where required.	True
donate	bool	If True (default), the input buffers may be reused by the native call and must not be used afterwards. Pass False to preserve them at the cost of additional memory.	True

Returns:

Type	Description
Union[Array, Tuple[Array, Array]]	Array or (Array, Array): The least-squares solution, and optionally the
Union[Array, Tuple[Array, Array]]	native per-rank diagnostic status.

Raises:

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

Notes

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

---

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

Solve a least-squares system inside a caller-owned jax.jit.

This lower-level interface performs the same work as :func:jaxmg.least_squares but leaves the outer JIT boundary to the caller. It returns a_work and b_work so an outer compiled function can donate both inputs into shape-compatible native work outputs.

Parameters:

Name	Type	Description	Default
a	Array	Rank-2 input matrix sharded over a one- or two-axis device mesh.	required
b	Array	Rank-1 or rank-2 solve input with M rows.	required
T_A	int	Square cuSOLVERMp tile width.	required
mesh	Mesh or AbstractMesh	JAX mesh used by jax.shard_map. If omitted, read from the sharding of a (its type inside jax.jit), or taken from the context mesh.	None
matrix_specs	PartitionSpec or tuple / list[PartitionSpec]	Rank-2 matrix sharding. If omitted, read from 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), add tile-aligned local capacity where required.	True

Returns:

Name	Type	Description
tuple	Array	(a_work, b_work, x, status) containing the opaque matrix and
	Array	solve-input work buffers, least-squares solution, and native per-rank
	Array	status. The first N rows of b_work contain x.

Raises:

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

Notes

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

---

Use least_squares_shardmap_ctx for local blocks inside your own jax.shard_map. The previous global-array context interface is now named least_squares_jit_ctx.

jaxmg.least_squares_shardmap_ctx(a, b, T_A, matrix_specs=None, *, in_specs=None, pad=True, _prepared=None)

Solve a least-squares system inside a caller's jax.shard_map.

A follows matrix_specs; B shares its row sharding and may have distributed or replicated columns. Replicated copies become duplicate right-hand sides in one distributed QR solve, not separate factorizations. Every process column must own a logical B tile, including these copies.

Parameters:

Name	Type	Description	Default
a	Array	Local block of a global M x N matrix, with M >= N.	required
b	Array	Local RHS matrix or vector, with A's dtype and row layout.	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 local tile-capacity padding; otherwise require aligned blocks.	True

Returns:

Type	Description
Array	(a_work, b_work, x, status). A work is padded and opaque; B work
Array	retains the logical M-row layout. The native backend restores the
Array	N-row solution across process rows, preserving B's column replication.
Array	Return both work buffers from an outer jitted function when donating inputs.
Tuple[Array, Array, Array, Array]	Status contains this rank's diagnostics.

Raises:

Type	Description
TypeError	If dtypes or partition specifications are unsupported.
ValueError	If called outside shard_map, shapes are incompatible, or a process column owns no logical B tile.