Add non-symmetric eigenvalue solvers, eig
/ eivals
, to the linalg extension
#935
Labels
RFC
Request for comments. Feature requests and proposed changes.
Currently the linalg extension only contains symmetric/hermitian eigenvalue routines,
eigh
andeigvalsh
. 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:
np.linalg.eig
,np.linalg.eigvals
are available: https://numpy.org/doc/2.1/reference/generated/numpy.linalg.eig.htmltorch.linalg.eig
,torch.linalg.eigvals
are available: https://pytorch.org/docs/stable/generated/torch.linalg.eig.htmljax.numpy.linalg.eig
,jax.numpy.linalg.eigvals
are available: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.linalg.eig.htmlUnder the hood, it uses the same MAGMA kernels as pytorch, so [1] applies. However, jax-ml/jax#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: cupy/cupy#8980, hopefully to CuPy 14.x.
Proposed API
The API should closely follow
eigh
andeigvalsh
: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 thateigenvalues[i]
corresponds toeigenvectors[:, 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] cupy/cupy#6359 (comment)
[2] https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html
The text was updated successfully, but these errors were encountered: