jaxmg.polar¤

polar computes the polar decomposition of a tall or square matrix \(A\):

\[ A = U_p H. \]

For a full-rank \(M \times N\) matrix with \(M \geq N\), the columns of \(U_p\) are orthonormal and \(H\) is an \(N \times N\) Hermitian positive-definite matrix. Set compute_h=False when only \(U_p\) is required; the \(H\) output and its reverse redistribution are then omitted.

Use polar for a direct decomposition. Use polar_shardmap_ctx when the decomposition is part of a larger caller-owned jax.jit. cuSOLVERMp overwrites the input matrix with \(U_p\), so the context interface returns \(U_p\) directly as the donated input alias rather than returning a separate work buffer.

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

Compute a distributed polar decomposition with cuSOLVERMp.

For a full-rank m x n matrix with m >= n, the decomposition is A = Up @ H. Up has orthonormal columns and H is an n x n Hermitian positive-definite matrix. Rank-deficient inputs instead produce a positive-semidefinite H and may not yield orthonormal columns in the null-space directions of Up.

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`
`compute_h`	`bool`	Whether to compute and return `H`. Default is True. This must be a Python `bool` fixed while tracing.	`True`
`return_status`	`bool`	If True, append the native per-rank diagnostic status vector to the 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 may be donated and overwritten by `Up`. Pass False to preserve it, at the cost of an additional A-sized allocation.	`True`

Returns:

Type	Description
`Array \| tuple[Array, ...]`	`(Up, H)` by default, or only `Up` when `compute_h=False`. If
`Array \| tuple[Array, ...]`	`return_status=True`, the status vector is appended to that result.
`Array \| tuple[Array, ...]`	Requested matrix outputs are restored to replicated sharding when
`Array \| tuple[Array, ...]`	`a` is replicated; otherwise they use `matrix_specs`.

Raises:

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

Notes

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

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

Compute a polar decomposition inside a caller-owned jax.jit.

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