diff --git a/.github/workflows/build-ci.yml b/.github/workflows/build-ci.yml index f782d76..18bc855 100644 --- a/.github/workflows/build-ci.yml +++ b/.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,6 +55,10 @@ jobs: tests: runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10", "3.13"] + steps: - name: Checkout repository uses: actions/checkout@v5 @@ -62,7 +66,7 @@ jobs: - 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 diff --git a/LICENSE b/LICENSE index cf472ee..271b2ae 100644 --- a/LICENSE +++ b/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: diff --git a/README.md b/README.md index 2f5491c..ff92733 100644 --- a/README.md +++ b/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 diff --git a/docs/speed.png b/docs/speed.png index 3475160..a870d79 100644 Binary files a/docs/speed.png and b/docs/speed.png differ diff --git a/notebooks/test.ipynb b/notebooks/test.ipynb index f263114..cbb3ebe 100644 --- a/notebooks/test.ipynb +++ b/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, diff --git a/pyproject.toml b/pyproject.toml index ed39d9c..c5edd6d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "jax>=0.4.36", "jaxlib", "numpy", + "packaging", "scipy", "torch", ] diff --git a/src/jeig/jeig.py b/src/jeig/jeig.py index 0d6fc90..2568ca3 100644 --- a/src/jeig/jeig.py +++ b/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,25 +107,57 @@ 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: @@ -120,11 +165,22 @@ def _eig_magma(matrix: jnp.ndarray, force_x64: bool) -> Tuple[jnp.ndarray, jnp.n 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, diff --git a/tests/test_eig.py b/tests/test_eig.py index d09c849..081fa66 100644 --- a/tests/test_eig.py +++ b/tests/test_eig.py @@ -1,4 +1,7 @@ -"""Tests for jax-wrapped eigendecomposition.""" +"""Tests for jax-wrapped eigendecomposition. + +Copyright (c) 2025 invrs.io LLC +""" import itertools import unittest @@ -15,15 +18,21 @@ BACKENDS = [ - jeig.EigBackend.JAX, + jeig.EigBackend.LAPACK, jeig.EigBackend.NUMPY, jeig.EigBackend.SCIPY, jeig.EigBackend.TORCH, ] + +# Only test the cusolver backend if supported by the installed jax version. +if _jeig._SUPPORTS_CUSOLVER and jax.devices()[0].platform == "gpu": + BACKENDS.append(jeig.EigBackend.CUSOLVER) + # Only test the magma backend if supported by the installed jax and torch versions. -if _jeig._JAX_HAS_MAGMA: +if _jeig._SUPPORTS_MAGMA: BACKENDS.append(jeig.EigBackend.MAGMA) + SHAPES = [(1, 2, 2), (1, 16, 16), (2, 16, 16), (2, 64, 64)] @@ -58,7 +67,9 @@ def test_backends_against_jax_symmetric_real(self, backend, shape): matrix = jax.random.normal(jax.random.PRNGKey(0), shape) matrix = 0.5 * (matrix + jnp.swapaxes(matrix, -1, -2)) - expected_eigval, expected_eigvec = jeig.eig(matrix, backend=jeig.EigBackend.JAX) + expected_eigval, expected_eigvec = jeig.eig( + matrix, backend=jeig.EigBackend.LAPACK + ) eigval, eigvec = jeig.eig(matrix, backend=backend) eigval, eigvec = _match_eigs(eigval, eigvec, expected_eigval, expected_eigvec) @@ -69,7 +80,9 @@ def test_backends_against_jax_symmetric_real(self, backend, shape): def test_backends_against_jax_real(self, backend, shape): matrix = jax.random.normal(jax.random.PRNGKey(0), shape) - expected_eigval, expected_eigvec = jeig.eig(matrix, backend=jeig.EigBackend.JAX) + expected_eigval, expected_eigvec = jeig.eig( + matrix, backend=jeig.EigBackend.LAPACK + ) eigval, eigvec = jeig.eig(matrix, backend=backend) eigval, eigvec = _match_eigs(eigval, eigvec, expected_eigval, expected_eigvec) @@ -81,7 +94,9 @@ def test_backends_against_jax_complex(self, backend, shape): real, imag = jax.random.normal(jax.random.PRNGKey(0), (2,) + shape) matrix = real + 1j * imag - expected_eigval, expected_eigvec = jeig.eig(matrix, backend=jeig.EigBackend.JAX) + expected_eigval, expected_eigvec = jeig.eig( + matrix, backend=jeig.EigBackend.LAPACK + ) eigval, eigvec = jeig.eig(matrix, backend=backend) eigval, eigvec = _match_eigs(eigval, eigvec, expected_eigval, expected_eigvec) @@ -96,7 +111,7 @@ def test_set_backend(self, backend): jeig.set_backend(jeig.EigBackend.TORCH) self.assertEqual(_jeig._DEFAULT_BACKEND, jeig.EigBackend.TORCH) - @parameterized.expand(["jax", "numpy", "scipy", "torch"]) + @parameterized.expand(["lapack", "numpy", "scipy", "torch"]) def test_set_backend_with_str(self, backend_str): self.assertEqual(_jeig._DEFAULT_BACKEND, jeig.EigBackend.TORCH) jeig.set_backend(backend_str)