jax.lax.linalg.eig#
- jax.lax.linalg.eig(x, *, compute_left_eigenvectors=True, compute_right_eigenvectors=True, use_magma=None)[source]#
Eigendecomposition of a general matrix.
Nonsymmetric eigendecomposition is only implemented on CPU and GPU. On GPU, the default implementation calls LAPACK directly on the host CPU, but an experimental GPU implementation using MAGMA is also available. The MAGMA implementation is typically slower than the equivalent LAPACK implementation for small matrices (less than about 2048), but it may perform better for larger matrices.
To enable the MAGMA implementation, you must install MAGMA yourself (there are Debian and conda-forge packages, or you can build from source). Then set the
use_magma
argument toTrue
, or set thejax_use_magma
configuration variable to"on"
or"auto"
:jax.config.update('jax_use_magma', 'on')
JAX will try to
dlopen
the installed MAGMA shared library, raising an error if it is not found. To explicitly specify the path to the MAGMA library, set the environment variable JAX_GPU_MAGMA_PATH to the full installation path.If
jax_use_magma
is set to"auto"
, the MAGMA implementation will be used if the library can be found, and the input matrix is sufficiently large (>= 2048x2048).- Parameters:
x (ArrayLike) – A batch of square matrices with shape
[..., n, n]
.compute_left_eigenvectors (bool) – If true, the left eigenvectors will be computed.
compute_right_eigenvectors (bool) – If true, the right eigenvectors will be computed.
use_magma (bool | None | None) – Locally override the
jax_use_magma
flag. IfTrue
, the eigendecomposition is computed using MAGMA. IfFalse
, the computation is done using LAPACK on to the host CPU. IfNone
(default), the behavior is controlled by thejax_use_magma
flag. This argument is only used on GPU.
- Returns:
The eigendecomposition of
x
, which is a tuple of the form(w, vl, vr)
wherew
are the eigenvalues,vl
are the left eigenvectors, andvr
are the right eigenvectors.vl
andvr
are optional and will only be included ifcompute_left_eigenvectors
orcompute_right_eigenvectors
respectively areTrue
.If the eigendecomposition fails, then arrays full of NaNs will be returned for that batch element.
- Return type: