Support cusolver backend

.github/workflows/build-ci.yml

@@ -18,7 +18,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.13"
           cache: "pip"
           cache-dependency-path: pyproject.toml
       - name: Test pre-commit hooks
@@ -36,7 +36,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.13"
           cache: "pip"
           cache-dependency-path: pyproject.toml
 
@@ -55,14 +55,18 @@ jobs:
 
   tests:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.13"]
+
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
 
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: ${{ matrix.python-version }}
           cache: "pip"
           cache-dependency-path: pyproject.toml
 

LICENSE

@@ -1,6 +1,6 @@
 BSD 3-Clause License
 
-Copyright (c) 2024, Martin Schubert
+Copyright (c) 2025 invrs.io LLC
 
 Redistribution and use in source and binary forms, with or without
 modification, are permitted provided that the following conditions are met:

README.md

@@ -1,10 +1,10 @@
 # jeig - Eigendecompositions wrapped for jax
-![Continuous integration](https://github.com/mfschubert/jeig/actions/workflows/build-ci.yml/badge.svg)
-![PyPI version](https://img.shields.io/pypi/v/jeig)
+[![Continuous integration](https://github.com/mfschubert/jeig/actions/workflows/build-ci.yml/badge.svg)](https://github.com/mfschubert/jeig/actions)
+[![PyPI version](https://img.shields.io/pypi/v/jeig)](https://pypi.org/project/jeig/)
 
 ## Overview
 
-This package wraps eigendecompositions as provided by jax, magma, numpy, scipy, and torch for use with jax. Depending upon your system and your versions of these packages, you may observe significant speed differences. The following were obtained using jax 0.4.37 on a system with 28-core Intel Xeon w7-3465X and NVIDIA RTX4090.
+This package wraps eigendecompositions as provided by jax, cusolver, magma, numpy, scipy, and torch for use with jax. Depending upon your system and your versions of these packages, you may observe significant speed differences. The following were obtained using jax 0.8.0 on a system with 28-core Intel Xeon w7-3465X and NVIDIA RTX4090.
 
 ![Speed comparison](https://github.com/mfschubert/jeig/blob/main/docs/speed.png?raw=true)
 
@@ -21,24 +21,35 @@ This will also install torch. If you only need torch for use with jeig, then the
 import jax
 import jeig
 
-matrix = jax.random.normal(jax.random.PRNGKey(0), (16, 1024, 1024))
-
-%timeit jax.block_until_ready(jeig.eig(matrix, backend="jax"))
+matrix = jax.random.normal(jax.random.PRNGKey(0), (1, 2048, 2048)).astype(complex)
 
+%timeit jax.block_until_ready(jeig.eig(matrix, backend="cusolver"))
+%timeit jax.block_until_ready(jeig.eig(matrix, backend="lapack"))
 %timeit jax.block_until_ready(jeig.eig(matrix, backend="magma"))
+%timeit jax.block_until_ready(jeig.eig(matrix, backend="torch"))
+```
+```
+1.31 s ± 43 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
+5.44 s ± 379 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
+11.1 s ± 937 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
+4.93 s ± 92.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
+```
 
-%timeit jax.block_until_ready(jeig.eig(matrix, backend="numpy"))
+The default torch backend has good performance when performing batched eigendecomposition on many-core CPUs.
 
-%timeit jax.block_until_ready(jeig.eig(matrix, backend="scipy"))
+```python
+matrix = jax.random.normal(jax.random.PRNGKey(0), (8, 2048, 2048)).astype(complex)
 
+%timeit jax.block_until_ready(jeig.eig(matrix, backend="cusolver"))
+%timeit jax.block_until_ready(jeig.eig(matrix, backend="lapack"))
+%timeit jax.block_until_ready(jeig.eig(matrix, backend="magma"))
 %timeit jax.block_until_ready(jeig.eig(matrix, backend="torch"))
 ```
 ```
-6.81 s ± 54.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
-1min 15s ± 1.35 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
-28.6 s ± 341 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
-14.8 s ± 396 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
-1.43 s ± 77.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
+10.4 s ± 116 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
+48.1 s ± 6.74 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
+1min 33s ± 1.49 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
+7.18 s ± 91.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
 ```
 
 ## Credit

notebooks/test.ipynb

@@ -22,10 +22,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "batch_size = [1, 16]\n",
+    "batch_size = [1, 8]\n",
     "matrix_size = [16, 32, 64, 128, 256, 512, 1024, 2048, 4096]\n",
     "repeats = 3\n",
-    "backends = [jeig.EigBackend.JAX, jeig.EigBackend.MAGMA, jeig.EigBackend.NUMPY, jeig.EigBackend.SCIPY, jeig.EigBackend.TORCH]\n",
+    "backends = [\n",
+    "    jeig.EigBackend.CUSOLVER,\n",
+    "    jeig.EigBackend.LAPACK,\n",
+    "    jeig.EigBackend.MAGMA,\n",
+    "    jeig.EigBackend.TORCH,\n",
+    "]\n",
     "\n",
     "fns = {}\n",
     "for backend in backends:\n",
@@ -41,7 +46,7 @@
     "        for backend in backends:\n",
     "            for repeat in range(repeats):\n",
     "                key = jax.random.fold_in(jax.random.PRNGKey(0), repeat)\n",
-    "                matrix = jax.random.normal(key, shape)\n",
+    "                matrix = jax.random.normal(key, shape).astype(complex)\n",
     "                t0 = time.time()\n",
     "                jax.block_until_ready(fns[backend](matrix))\n",
     "                et = time.time() - t0\n",
@@ -69,11 +74,45 @@
     "    ax.set_xlabel(\"Matrix size\")\n",
     "plt.tight_layout()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "import jeig\n",
+    "\n",
+    "matrix = jax.random.normal(jax.random.PRNGKey(0), (8, 2048, 2048)).astype(complex)\n",
+    "\n",
+    "%timeit jax.block_until_ready(jeig.eig(matrix, backend=\"cusolver\"))\n",
+    "%timeit jax.block_until_ready(jeig.eig(matrix, backend=\"lapack\"))\n",
+    "%timeit jax.block_until_ready(jeig.eig(matrix, backend=\"magma\"))\n",
+    "%timeit jax.block_until_ready(jeig.eig(matrix, backend=\"torch\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "import jeig\n",
+    "\n",
+    "matrix = jax.random.normal(jax.random.PRNGKey(0), (1, 2048, 2048)).astype(complex)\n",
+    "\n",
+    "%timeit jax.block_until_ready(jeig.eig(matrix, backend=\"cusolver\"))\n",
+    "%timeit jax.block_until_ready(jeig.eig(matrix, backend=\"lapack\"))\n",
+    "%timeit jax.block_until_ready(jeig.eig(matrix, backend=\"magma\"))\n",
+    "%timeit jax.block_until_ready(jeig.eig(matrix, backend=\"torch\"))"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "invrs2",
+   "display_name": "invrs5",
    "language": "python",
    "name": "python3"
   },
@@ -87,7 +126,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.13.9"
   }
  },
  "nbformat": 4,

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "jax>=0.4.36",
     "jaxlib",
     "numpy",
+    "packaging",
     "scipy",
     "torch",
 ]

src/jeig/jeig.py

@@ -1,10 +1,14 @@
-"""Various implementations of `eig` wrapped for use with jax."""
+"""Various implementations of `eig` wrapped for use with jax.
+
+Copyright (c) 2025 invrs.io LLC
+"""
 
 import enum
 import functools
 import multiprocessing as mp
 import os
 import warnings
+from packaging import version
 from typing import Any, List, Optional, Tuple
 
 import jax
@@ -25,7 +29,14 @@
         os.path.dirname(torch.__file__), "lib", "libtorch_cuda_linalg.so"
     )
 
-_JAX_HAS_MAGMA = torch.cuda.has_magma
+# Versions of jax newer than `0.8.0` have an `jax.lax.linalg.eig` function with an
+# `implementation` argument, and support eigendecomposition via cusolver (if on GPU).
+_SUPPORTS_IMPLEMENTATION = version.parse(jax.__version__) >= version.parse("0.8.0")
+_SUPPORTS_CUSOLVER = _SUPPORTS_IMPLEMENTATION and jax.devices()[0].platform == "gpu"
+
+# Identify whether the magma backend is available.
+_SUPPORTS_MAGMA = torch.cuda.has_magma and jax.devices()[0].platform == "gpu"
+
 
 callback = functools.partial(jax.pure_callback, vmap_method="expand_dims")
 callback_sequential = functools.partial(jax.pure_callback, vmap_method="sequential")
@@ -36,7 +47,8 @@
 
 @enum.unique
 class EigBackend(enum.Enum):
-    JAX = "jax"
+    CUSOLVER = "cusolver"
+    LAPACK = "lapack"
     MAGMA = "magma"
     NUMPY = "numpy"
     SCIPY = "scipy"
@@ -46,7 +58,8 @@ class EigBackend(enum.Enum):
     def from_string(cls, backend: str) -> "EigBackend":
         """Returns the specified backend."""
         return {
-            cls.JAX.value: cls.JAX,
+            cls.CUSOLVER.value: cls.CUSOLVER,
+            cls.LAPACK.value: cls.LAPACK,
             cls.MAGMA.value: cls.MAGMA,
             cls.NUMPY.value: cls.NUMPY,
             cls.SCIPY.value: cls.SCIPY,
@@ -94,37 +107,80 @@ def eig(
     return EIG_FNS[backend](matrix, force_x64)
 
 
-def _eig_jax(matrix: jnp.ndarray, force_x64: bool) -> Tuple[jnp.ndarray, jnp.ndarray]:
-    """Eigendecomposition using `jax.numpy.linalg.eig`."""
+def _eig_cusolver(
+    matrix: jnp.ndarray, force_x64: bool
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Eigendecomposition using `jax.lax.linalg.eig` with `cusolver` backend."""
+    if not _SUPPORTS_CUSOLVER:
+        raise RuntimeError("`CUSOLVER` backend is not available.")
     if force_x64:
         matrix_dtype = jnp.promote_types(matrix.dtype, jnp.float64)
     else:
         matrix_dtype = matrix.dtype
     output_dtype = jnp.promote_types(matrix, jnp.complex64)
+
     eigval, eigvec = jax.lax.linalg.eig(
         matrix.astype(matrix_dtype),
         compute_left_eigenvectors=False,
-        use_magma=False,
+        implementation=jax.lax.linalg.EigImplementation.CUSOLVER,
     )
     return eigval.astype(output_dtype), eigvec.astype(output_dtype)
 
 
+def _eig_lapack(
+    matrix: jnp.ndarray, force_x64: bool
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Eigendecomposition using `jax.lax.linalg.eig` with `lapack` backend."""
+    if force_x64:
+        matrix_dtype = jnp.promote_types(matrix.dtype, jnp.float64)
+    else:
+        matrix_dtype = matrix.dtype
+    output_dtype = jnp.promote_types(matrix, jnp.complex64)
+
+    if _SUPPORTS_IMPLEMENTATION:
+        eig_fn = functools.partial(
+            jax.lax.linalg.eig,
+            compute_left_eigenvectors=False,
+            implementation=jax.lax.linalg.EigImplementation.LAPACK,
+        )
+    else:
+        eig_fn = functools.partial(
+            jax.lax.linalg.eig,
+            compute_left_eigenvectors=False,
+            use_magma=False,
+        )
+
+    eigval, eigvec = eig_fn(matrix.astype(matrix_dtype))
+    return eigval.astype(output_dtype), eigvec.astype(output_dtype)
+
+
 def _eig_magma(matrix: jnp.ndarray, force_x64: bool) -> Tuple[jnp.ndarray, jnp.ndarray]:
-    """Eigendecomposition using `jax.numpy.linalg.eig`."""
-    if not _JAX_HAS_MAGMA:
-        raise ValueError(
+    """Eigendecomposition using `jax.lax.linalg.eig` with `magma` backend."""
+    if not _SUPPORTS_MAGMA:
+        raise RuntimeError(
             "`MAGMA` backend is not available; `torch.cuda.has_magma` is `False`."
         )
     if force_x64:
         matrix_dtype = jnp.promote_types(matrix.dtype, jnp.float64)
     else:
         matrix_dtype = matrix.dtype
     output_dtype = jnp.promote_types(matrix, jnp.complex64)
-    eigval, eigvec = jax.lax.linalg.eig(
-        matrix.astype(matrix_dtype),
-        compute_left_eigenvectors=False,
-        use_magma=True,
-    )
+
+    eig_fn = functools.partial(jax.lax.linalg.eig, compute_left_eigenvectors=False)
+    if _SUPPORTS_IMPLEMENTATION:
+        eig_fn = functools.partial(
+            jax.lax.linalg.eig,
+            compute_left_eigenvectors=False,
+            implementation=jax.lax.linalg.EigImplementation.MAGMA,
+        )
+    else:
+        eig_fn = functools.partial(
+            jax.lax.linalg.eig,
+            compute_left_eigenvectors=False,
+            use_magma=True,
+        )
+
+    eigval, eigvec = eig_fn(matrix.astype(matrix_dtype))
     return eigval.astype(output_dtype), eigvec.astype(output_dtype)
 
 
@@ -185,7 +241,7 @@ def _eig_fn(m: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
 
 
 def _eig_torch(matrix: jnp.ndarray, force_x64: bool) -> Tuple[jnp.ndarray, jnp.ndarray]:
-    """Eigendecomposition using `torc.linalg.eig`."""
+    """Eigendecomposition using `torch.linalg.eig`."""
     dtype = jnp.promote_types(matrix.dtype, jnp.complex64)
 
     def _eig_fn(matrix: jnp.ndarray) -> Tuple[NDArray, NDArray]:
@@ -228,7 +284,8 @@ def _eig_torch_parallelized(x: torch.Tensor) -> List[Tuple[torch.Tensor, torch.T
 
 
 EIG_FNS = {
-    EigBackend.JAX: _eig_jax,
+    EigBackend.CUSOLVER: _eig_cusolver,
+    EigBackend.LAPACK: _eig_lapack,
     EigBackend.MAGMA: _eig_magma,
     EigBackend.NUMPY: _eig_numpy,
     EigBackend.SCIPY: _eig_scipy,