Enable optional use of jeig for eigendecomposition (#127)

* Use jeig for eigendecomposition

* make jeig dep optional

* make jeig dep optional

* tests with jeig

* make jeig dep optional

* reorder imports

* install jeig for lint and typecheck

---------

Co-authored-by: Martin Schubert <[email protected]>

.github/workflows/build-ci.yml

@@ -23,7 +23,7 @@ jobs:
       - name: Setup environment
         run: |
             python -m pip install --upgrade pip
-            pip install ".[dev]"
+            pip install ".[dev,jeig]"
       - name: Lint Python files
         run: |
             find . -name "*.py" | xargs black --check
@@ -55,6 +55,24 @@ jobs:
       - name: Test fmmax
         run: pytest tests/fmmax
 
+  test-fmmax-jeig:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      - name: Setup python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+          cache-dependency-path: pyproject.toml
+      - name: Setup environment
+        run: |
+            python -m pip install --upgrade pip
+            pip install ".[tests,dev,jeig]"        
+      - name: Test fmmax
+        run: pytest tests/fmmax
+
   test-grcwa:
     runs-on: ubuntu-latest
     steps:

pyproject.toml

@@ -25,8 +25,10 @@ dependencies = [
   "jaxlib",
   "numpy",
 ]
-
 [project.optional-dependencies]
+jeig = [
+  "jeig",
+]
 tests = [
   "grcwa",
   "parameterized",

src/fmmax/utils.py

@@ -8,6 +8,18 @@
 import jax
 import jax.numpy as jnp
 
+# The `jeig` package offers several jax-wrapped implementations of eigendecomposition,
+# some of which have performance benefits. However, since `jeig` has a dependency on
+# pytorch, we make its use optional. If `jeig` is not available, we fall back on a
+# pure-jax implementation of the eigendecomposition.
+try:
+    import jeig
+
+    _JEIG_AVALABLE = True
+except ModuleNotFoundError:
+    _JEIG_AVALABLE = False
+
+
 EIG_EPS_RELATIVE = 1e-12
 EIG_EPS_MINIMUM = 1e-24
 
@@ -120,10 +132,10 @@ def eig(
         The eigenvalues and eigenvectors.
     """
     del eps_relative
-    return _eig_host(matrix)
+    return _eig(matrix)
 
 
-def _eig_host(matrix: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
+def _eig_host_jax(matrix: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
     """Wraps jnp.linalg.eig so that it can be jit-ed on a machine with GPUs."""
 
     def _eig_cpu(matrix: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
@@ -143,12 +155,20 @@ def _eig_cpu(matrix: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
     )
 
 
+def _eig(matrix: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Eigendecomposition using `jeig` if available, and `_eig_host_jax` if not."""
+    if _JEIG_AVALABLE:
+        return jeig.eig(matrix)
+    else:
+        return _eig_host_jax(matrix)
+
+
 def _eig_fwd(
     matrix: jnp.ndarray,
     eps_relative: float,
 ) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, jnp.ndarray, float]]:
     """Implements the forward calculation for `eig`."""
-    eigenvalues, eigenvectors = _eig_host(matrix)
+    eigenvalues, eigenvectors = _eig(matrix)
     return (eigenvalues, eigenvectors), (eigenvalues, eigenvectors, eps_relative)
 
 

tests/fmmax/test_utils.py

@@ -146,8 +146,8 @@ def test_value_matches_eig_with_nondegenerate_eigenvalues(self):
             jax.device_put(matrix, device=jax.devices("cpu")[0])
         )
         eigval, eigvec = utils.eig(matrix)
-        onp.testing.assert_array_equal(eigval, expected_eigval)
-        onp.testing.assert_array_equal(eigvec, expected_eigvec)
+        onp.testing.assert_allclose(eigval, expected_eigval, rtol=1e-12)
+        onp.testing.assert_allclose(eigvec, expected_eigvec, rtol=1e-12)
 
     def test_eigvalue_jacobian_matches_expected_real_matrix(self):
         matrix = jax.random.normal(jax.random.PRNGKey(0), (2, 4, 4)).astype(complex)