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_shardmap_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_shardmap_ctx(a, b, T_A, mesh=None, matrix_specs=None, *, in_specs=None, pad=True)` ¤

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.