jaxmg.syevd¤

syevd computes the eigenvalues and, by default, eigenvectors of a symmetric real matrix or Hermitian complex matrix. Eigenvalues are real, while eigenvectors use the input dtype and are returned in the original JAX-facing matrix layout.

Set return_eigenvectors=False when only the spectrum is required. This selects the eigenvalues-only mode and reduces memory use by avoiding the full distributed eigenvector matrix.

Use syevd for a direct eigensolve. Use syevd_shardmap_ctx when the eigensolve is part of a larger caller-owned jax.jit. The context interface returns the native input-matrix work buffer so the outer compiled function can preserve the donated input alias. cuSOLVERMp stores the overwritten input matrix and eigenvectors in separate distributed buffers, so a_work is the alias target rather than the eigenvector output when vectors are requested.

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

Compute eigenvalues and optionally eigenvectors using multi-GPU SYEVD.

This is the high-level JAXMg symmetric/Hermitian eigensolver entry point. It accepts a block-sharded JAX matrix, prepares the tile-aligned local capacity required by cuSOLVERMp, and calls the fused native backend. By default it returns eigenvalues together with eigenvectors in the same JAX-facing matrix layout as the input. Setting return_eigenvectors=False selects cuSOLVERMp's eigenvalues-only mode and avoids allocating or restoring a matrix-sized eigenvector result.

Note

If a local shard dimension is not divisible by T_A, pad=True allocates additional tile-aligned capacity before the native call. Choosing a tile size that divides the local dimensions avoids this allocation. Performance depends on the matrix size, process grid, and tile size.

Parameters:

Name	Type	Description	Default
`a`	`Array`	A 2D symmetric/Hermitian matrix. Expected to be sharded across a 2D mesh with a matrix `PartitionSpec` such as `P(<row_axis>, <col_axis>)`.	required
`T_A`	`int`	Square tile width used by cuSOLVERMp. Each local shard dimension must be a multiple of `T_A` after padding.	required
`mesh`	`Mesh`	JAX mesh used for `jax.shard_map`. If omitted, inferred from `a.sharding.mesh`.	`None`
`matrix_specs`	`PartitionSpec or tuple / list[PartitionSpec]`	PartitionSpec describing the 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`
`return_eigenvectors`	`bool`	If True (default), compute and return eigenvectors as well as eigenvalues. If False, return only eigenvalues and use the shorter native values-only workflow. This must be a Python `bool` fixed while tracing.	`True`
`return_status`	`bool`	If True append a native per-rank diagnostic status vector to the return values. Default is False.	`False`
`pad`	`bool`	If True (default) apply per-device padding to meet `T_A` requirements; if False the caller must supply already correct shapes.	`True`

Returns:

Type	Description
`Array \| Tuple[Array, Array] \| Tuple[Array, Array, Array]`	If `return_eigenvectors=True`, `(eigenvalues, eigenvectors)` or
`Array \| Tuple[Array, Array] \| Tuple[Array, Array, Array]`	`(eigenvalues, eigenvectors, status)`. If False, `eigenvalues` or
`Array \| Tuple[Array, Array] \| Tuple[Array, Array, Array]`	`(eigenvalues, status)`.

Raises:

Type	Description
`TypeError`	If dtypes or `PartitionSpec` inputs are unsupported.
`ValueError`	If shapes, tile sizes, or mesh layouts are incompatible.

Notes

The FFI call can donate the input buffer to enable zero-copy interaction with the native library.
Native code converts row-major JAX local storage to cuSOLVERMp's column-major local layout and redistributes to 2D block-cyclic layout. The eigenvector mode also redistributes its matrix result back to the original JAX layout.
If the native solver fails the outputs may contain NaNs and the status, when requested, will be non-zero.

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

Compute eigenvalues and optional eigenvectors while exposing work storage.

This helper is the lower-level variant of :func:jaxmg.syevd intended for contexts where the caller controls the outer jax.jit boundary. It performs the same validation, local padding, shard-map construction, and fused cuSOLVERMp FFI call as the public eigensolver, but does not introduce an internal jax.jit. Instead, it returns the native matrix work buffer so an outer JIT can donate a into an A-sized output.

cuSOLVERMp uses separate distributed matrices for its overwritten input d_A and eigenvector output d_Q. The context interface therefore returns a_work to provide the outer compiled function with an A-sized alias target for the donated input. When eigenvectors are requested, they occupy a separate matrix-sized allocation. Values-only execution omits that allocation and the associated reverse redistribution.