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

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.

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
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	(Up, H, status) when compute_h=True or (Up, status)
...	otherwise. The status vector is always returned so an enclosing
tuple[Array, ...]	compiled function can propagate native diagnostics. Requested matrix
tuple[Array, ...]	outputs are restored when a is replicated; otherwise they use
tuple[Array, ...]	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.

---

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

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

Compute polar factors inside a caller's jax.shard_map.

The enclosing mesh determines the global matrix dimensions; only the global matrix, not each local block, must be tall or square.

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
compute_h	bool	Static Python bool selecting the square H output.	True
pad	bool	Allow tile-capacity padding; otherwise require aligned blocks.	True

Returns:

Name	Type	Description
tuple	tuple[Array, ...]	(Up, H, status) or (Up, status) when compute_h=False. Factors are local blocks with padding removed; status contains this rank's diagnostics. Up aliases the input work buffer.

Raises:

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