jaxmg.qr¤

qr computes the reduced QR decomposition of a tall or square matrix \(A\):

\[ A = Q R, \qquad Q^{\dagger}Q=I, \]

where an \(M\times N\) input with \(M\geq N\) produces an \(M\times N\) matrix \(Q\) and an \(N\times N\) upper-triangular matrix \(R\).

Use qr for a direct decomposition. Use qr_shardmap_ctx when the decomposition is part of a larger caller-owned jax.jit. cuSOLVERMp overwrites the input matrix with \(Q\), so the context interface returns \(Q\) directly as the donated input alias.

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

Compute a distributed reduced QR decomposition with cuSOLVERMp.

For a tall or square m x n matrix with m >= n, returns an m x n matrix Q with orthonormal columns and an n x n upper-triangular matrix R such that A = Q @ R.

Parameters:

Name	Type	Description	Default
`a`	`Array`	A rank-2 real or complex tall or square matrix sharded over a one- or two-axis device mesh.	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 the native per-rank diagnostic status vector after `Q` and `R`.	`False`
`pad`	`bool`	If True (default), add tile-aligned local capacity where required. If False, all participating local matrix shapes must already be divisible by `T_A`.	`True`
`donate`	`bool`	If True (default), the input may be donated and overwritten by `Q`. Pass False to preserve it, at the cost of an additional A-sized allocation.	`True`

Returns:

Type	Description
`Array`	`(Q, R)` by default, or `(Q, R, status)` when
`...`	`return_status=True`. `Q` and `R` are restored to replicated
`tuple[Array, ...]`	sharding when `a` is replicated; otherwise they use `matrix_specs`.

Raises:

Type	Description
`TypeError`	If the input dtype or sharding specification is unsupported.
`ValueError`	If the matrix is wide or its shape, tile size, process grid, or output layout is incompatible with cuSOLVERMp.

Notes

If the native solver fails, Q and R are filled with NaN and the status is non-zero.

`jaxmg.qr_shardmap_ctx(a, T_A, mesh=None, matrix_specs=None, *, in_specs=None, pad=True)` ¤

Compute reduced QR inside a caller-owned jax.jit.

This interface performs the same validation, redistribution, and native computation as :func:jaxmg.qr, but leaves the outer compilation and donation boundary to the caller. The native Q output aliases the input work buffer before any final restoration of replicated sharding.