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 to True, or set the jax_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) – Locally override the jax_use_magma flag. If True, the eigendecomposition is computed using MAGMA. If False, the computation is done using LAPACK on to the host CPU. If None (default), the behavior is controlled by the jax_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) where w are the eigenvalues, vl are the left eigenvectors, and vr are the right eigenvectors. vl and vr are optional and will only be included if compute_left_eigenvectors or compute_right_eigenvectors respectively are True.

If the eigendecomposition fails, then arrays full of NaNs will be returned for that batch element.

Return type:

list[Array]