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_jit_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.

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

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_jit_ctx(a, T_A, mesh=None, matrix_specs=None, *, in_specs=None, pad=True, _prepared=None)

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.

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
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

Returns:

Type	Description
Array	(Q, R, status). The status vector is always returned so an
Array	enclosing compiled function can propagate native diagnostics. Q
Array	and R are restored when a is replicated; otherwise they use
tuple[Array, Array, Array]	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, matrix_specs=None, *, in_specs=None, pad=True)

Compute reduced QR factors inside a caller's jax.shard_map.

The global matrix must be tall or square. Local blocks need not satisfy that aspect ratio: the enclosing mesh determines the global dimensions.

Parameters:

Name	Type	Description	Default
a	Array	Local block of the global M x N matrix, with M >= N.	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	(q, r, status): local blocks of the M x N and N x N factors,
Array	respectively, and this rank's diagnostic vector. Q reuses the input
Array	work buffer, so no separate a_work output is needed for donation.

Raises:

Type	Description
TypeError	If the dtype or partition specification is unsupported.
ValueError	If called outside shard_map or input/output layouts are invalid.