array-api icon indicating copy to clipboard operation
array-api copied to clipboard

Add non-symmetric eigenvalue solvers, `eig` / `eivals`, to the linalg extension

Open ev-br opened this issue 7 months ago • 3 comments

Currently the linalg extension only contains symmetric/hermitian eigenvalue routines, eigh and eigvalsh. The general/non-symmetric versions are missing. IIUC one reason for the omission is that the low-level routines were not universally available on CUDA or were of allegedly questionable quality [1]. Now that the needed LAPACK-like routine was added to cuSolver 12.6 update 2 [2], the low-level blocker is no more.

Availability:

  • numpy

np.linalg.eig, np.linalg.eigvals are available: https://numpy.org/doc/2.1/reference/generated/numpy.linalg.eig.html

  • torch

torch.linalg.eig, torch.linalg.eigvals are available: https://pytorch.org/docs/stable/generated/torch.linalg.eig.html

  • jax

jax.numpy.linalg.eig, jax.numpy.linalg.eigvals are available: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.linalg.eig.html

Under the hood, it uses the same MAGMA kernels as pytorch, so [1] applies. However, https://github.com/jax-ml/jax/issues/27265 tracks a potential update to use the new routines from [2].

  • cupy

Not available in CuPy 13.4 yet; There is a pull request to add it: https://github.com/cupy/cupy/pull/8980, hopefully to CuPy 14.x.

Proposed API

The API should closely follow eigh and eigvalsh:

  • eigvals(x, /) should return an array of eigenvalues in an unspecified order;

  • eig(x, /) should return a namedtuple (eigenvalues, right eigenvectors). The order of eigenvalues and eigenvectors is unspecified, and the only requirement is that eigenvalues[i] corresponds to eigenvectors[:, i].

Whether the eigenvectors have a unit norm is unspecified and thus implementation-defined.

An option to return left eigenvectors may be added in the future if the array libraries add this option (currently, neither numpy nor torch provide this option).

[1] https://github.com/cupy/cupy/issues/6359#issuecomment-1019917948 [2] https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html

ev-br avatar Apr 29 '25 10:04 ev-br

Thanks @ev-br. This was discussed in the community meeting yesterday, and it seemed like a potentially good idea to the attendees. We noticed one hiccup: PyTorch and JAX always return arrays with complex dtype, while NumPy has value-dependent behavior (return dtype is float if imaginary values are all zero). It'd need checking if it's possible to change/fix this in NumPy - which may be tricky. Could you look into this?

rgommers avatar May 16 '25 05:05 rgommers

Here's a numpy issue to discuss this change: https://github.com/numpy/numpy/issues/29000. Also https://github.com/cupy/cupy/pull/8980#issuecomment-2890821763 is an attempt to convince CuPy devs to not follow numpy and have a dtype-stable API.

Overall, I believe that even if numpy keeps doing what it is doing, having eig in the standard with an implementation-defined return dtype is strictly better than not having it at all.

ev-br avatar May 20 '25 17:05 ev-br

CuPy merged the support for eig and eigvals in https://github.com/cupy/cupy/pull/8980/ for v14.0.a2. IIRC, the CuPy interface follows pytorch and jax.numpy and always return eigenvalues as a complex array (with zero imaginary parts if they happen to be on the real line).

ev-br avatar May 29 '25 08:05 ev-br