How to use jax.numpy.linalg.eigvals on GPU?

[shailesh1729]

Consider the following code:

import jax.numpy as jnp
I = jnp.eye(4)
jnp.linalg.eigvals(I)

It works on CPU but fails with the following error on GPU:

NotImplementedError: Nonsymmetric eigendecomposition is only implemented on the CPU backend

The documentation doesn't provide any details about how to indicate that the argument to the eigvals call is indeed symmetric.

What am I doing wrong?

Or is that the eigvals and eig are not implemented on GPU at all? The error message certainly doesn't convey that.

[jakevdp]

Indeed, eig and eigvals are not implemented on GPU, but I believe eigh and eigvalsh are (the h indicates that these routines are designed for hermitian matrices, and for real-valued inputs hermitian implies symmetric). If you are doing symmetric eigenvalue decomposition on any backend, you should prefer eigh or eigvalsh.