jaxmg.gesvd¤

gesvd computes the singular-value decomposition of an \(M \times N\) real or complex matrix \(A\):

\[ A = U \Sigma V^{\dagger}. \]

By default, gesvd returns the reduced decomposition (U, s, Vh), following jax.numpy.linalg.svd. Set full_matrices=True to return full-sized singular-vector matrices. The left and right singular vectors can be requested independently with compute_u and compute_vh; outputs that are not requested are neither allocated nor redistributed.

Use gesvd for a direct decomposition. Use gesvd_shardmap_ctx when the SVD is part of a larger caller-owned jax.jit. The context interface returns the overwritten input-matrix work buffer so an outer compiled function can preserve the donated input alias. cuSOLVERMp requires separate storage for A, U, and Vh, so requested singular-vector matrices cannot alias the donated input.

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

Compute a distributed singular-value decomposition with cuSOLVERMp.

This is the high-level JAXMg interface for rectangular or square matrices. For an m x n input, let k = min(m, n). By default the routine returns the reduced decomposition (U, s, Vh) with shapes (m, k), (k,), and (k, n). full_matrices=True instead requests U with shape (m, m) and Vh with shape (n, n). Either vector matrix can be disabled independently so JAX does not allocate or redistribute an output that the application does not require.

The input and every requested matrix output use the same JAX mesh and PartitionSpec. Their logical dimensions must therefore be divisible by the corresponding process-grid dimensions. Tile padding is applied separately to A, U, and Vh when their local dimensions are not divisible by T_A.

Parameters:

Name	Type	Description	Default
`a`	`Array`	A rank-2 real or complex matrix sharded over a one- or two-axis device mesh.	required
`T_A`	`int`	Square cuSOLVERMp tile width. GESVD supports rectangular matrices but requires equal row and column tile dimensions.	required
`mesh`	`Mesh`	JAX mesh used by `jax.shard_map`. If omitted, inferred from `a.sharding.mesh`.	`None`
`matrix_specs`	`PartitionSpec or tuple / list[PartitionSpec]`	Rank-2 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`
`compute_u`	`bool`	Whether to compute and return left singular vectors. Default is True. This must be a Python `bool` fixed while tracing.	`True`
`compute_vh`	`bool`	Whether to compute and return conjugate- transposed right singular vectors. Default is True. This must be a Python `bool` fixed while tracing.	`True`
`full_matrices`	`bool`	If False (default), return reduced vector matrices. If True, return full U and Vh matrices. This must be fixed while tracing and affects only requested vector outputs.	`False`
`return_status`	`bool`	If True, append the native per-rank diagnostic status vector to the selected numerical outputs.	`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 buffers may be donated to the native call for zero-copy execution, which means they are deleted and cannot be used again. Pass False to preserve them, at the cost of keeping the original and working buffers in memory simultaneously.	`True`

Returns:

Type	Description
`Array \| tuple[Array, ...]`	The selected numerical outputs follow NumPy/JAX SVD order:
`Array \| tuple[Array, ...]`	`(U, s, Vh)` when both vectors are requested, `(U, s)` for U only,
`Array \| tuple[Array, ...]`	`(s, Vh)` for Vh only, and `s` for values only. If
`Array \| tuple[Array, ...]`	`return_status=True`, the status vector is appended to that result.

Raises:

Type	Description
`TypeError`	If the dtype, static mode flags, or sharding specification is unsupported.
`ValueError`	If a shape, tile size, process grid, or requested output layout is incompatible with cuSOLVERMp.

Notes

The internally jitted implementation donates a to the opaque matrix work result returned by the native call.
cuSOLVERMp requires A, U, and Vh to occupy distinct storage. Donation therefore removes a second A-sized work allocation but cannot alias A to either singular-vector output.
If the native solver fails, numerical outputs may be incomplete; use return_status=True when per-rank diagnostics are required.

`jaxmg.gesvd_shardmap_ctx(a, T_A, mesh=None, matrix_specs=None, *, in_specs=None, compute_u=True, compute_vh=True, full_matrices=False, pad=True)` ¤

Compute a distributed SVD while exposing donated matrix work storage.

This lower-level interface performs the same validation, padding, native redistribution, and cuSOLVERMp execution as :func:jaxmg.gesvd, but leaves the outer jax.jit boundary to the caller. The first return value is the opaque a_work buffer overwritten by GESVD. Keeping that value in the outer function's returned pytree allows jax.jit(..., donate_argnums=(0,)) to alias the input matrix to an A-sized output.

Singular values and requested vector matrices follow the same shapes and selection rules as :func:jaxmg.gesvd. Unlike POTRS and LU solve, the factorized A buffer is not a numerical result. cuSOLVERMp also prohibits A, U, and Vh from overlapping, so requested vector matrices remain separate allocations even when A is donated.