From 656afae733ea0cb3fe650d2f568de592d8c950e1 Mon Sep 17 00:00:00 2001
From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com>
Date: Wed, 15 Jan 2025 13:54:34 +0000
Subject: [PATCH 1/3] make cueq optional dep and add special test

---
 .github/workflows/unittest.yaml | 33 +++++++++++++++++++++++++++------
 setup.cfg                       |  2 +-
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/unittest.yaml b/.github/workflows/unittest.yaml
index 857bf894..63b3da21 100644
--- a/.github/workflows/unittest.yaml
+++ b/.github/workflows/unittest.yaml
@@ -5,9 +5,8 @@ on:
     branches: [main]
 
 jobs:
-  pytest-container:
+  pytest-general:
     runs-on: ubuntu-latest
-
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
@@ -15,15 +14,37 @@ jobs:
           python-version: "3.10"
           cache: "pip"
 
-      - name: Install requirements
+      - name: Install requirements (general tests)
         run: |
           pip install -U pip
           pip install .[dev]
 
-      - name: Log installed environment
+      - name: Log installed environment (general tests)
+        run: |
+          python3 -m pip freeze
+
+      - name: Run general unit tests
+        run: |
+          pytest tests --ignore=tests/test_cueq.py
+
+  pytest-cueq:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+
+      - name: Install requirements (with cueq)
+        run: |
+          pip install -U pip
+          pip install .[dev, cueq]
+
+      - name: Log installed environment (with cueq)
         run: |
           python3 -m pip freeze
 
-      - name: Run unit tests
+      - name: Run cueq-specific tests
         run: |
-          pytest tests
+          pytest tests/test_cueq.py
diff --git a/setup.cfg b/setup.cfg
index be4804c4..0189c000 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -29,7 +29,6 @@ install_requires =
     GitPython
     pyYAML
     tqdm
-    cuequivariance-torch
     # for plotting:
     matplotlib
     pandas
@@ -60,5 +59,6 @@ dev =
     pytest-benchmark
     pylint
 schedulefree = schedulefree
+cueq = cuequivariance-torch
 cueq-cuda-11 = cuequivariance-ops-torch-cu11
 cueq-cuda-12 = cuequivariance-ops-torch-cu12

From 08039971c72791ed3bb80f463c37ce1f38fab21b Mon Sep 17 00:00:00 2001
From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com>
Date: Wed, 15 Jan 2025 14:05:01 +0000
Subject: [PATCH 2/3] update pytest-cueq install

---
 .github/workflows/unittest.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/unittest.yaml b/.github/workflows/unittest.yaml
index 63b3da21..68febfee 100644
--- a/.github/workflows/unittest.yaml
+++ b/.github/workflows/unittest.yaml
@@ -39,7 +39,7 @@ jobs:
       - name: Install requirements (with cueq)
         run: |
           pip install -U pip
-          pip install .[dev, cueq]
+          pip install ".[dev, cueq]"
 
       - name: Log installed environment (with cueq)
         run: |

From 140d25032f0229ed6d4420beb19f6f42d99c7fa8 Mon Sep 17 00:00:00 2001
From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com>
Date: Wed, 15 Jan 2025 14:43:35 +0000
Subject: [PATCH 3/3] move tolerance cueq

---
 .github/workflows/unittest.yaml | 2 +-
 tests/test_calculator.py        | 1 +
 tests/test_cueq.py              | 4 ++--
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/unittest.yaml b/.github/workflows/unittest.yaml
index 68febfee..1ee0ceeb 100644
--- a/.github/workflows/unittest.yaml
+++ b/.github/workflows/unittest.yaml
@@ -47,4 +47,4 @@ jobs:
 
       - name: Run cueq-specific tests
         run: |
-          pytest tests/test_cueq.py
+          pytest tests/test_cueq.py tests/test_calculator.py
diff --git a/tests/test_calculator.py b/tests/test_calculator.py
index 6c9e2563..b23ce709 100644
--- a/tests/test_calculator.py
+++ b/tests/test_calculator.py
@@ -586,6 +586,7 @@ def test_calculator_descriptor(fitting_configs, trained_equivariant_model):
     assert not np.allclose(desc, desc_rotated, atol=1e-6)
 
 
+@pytest.mark.skipif(not CUET_AVAILABLE, reason="cuequivariance not installed")
 def test_calculator_descriptor_cueq(fitting_configs, trained_equivariant_model_cueq):
     at = fitting_configs[2].copy()
     at_rotated = fitting_configs[2].copy()
diff --git a/tests/test_cueq.py b/tests/test_cueq.py
index 8d713c78..afba3311 100644
--- a/tests/test_cueq.py
+++ b/tests/test_cueq.py
@@ -143,7 +143,7 @@ def test_bidirectional_conversion(
         loss_e3nn_back.backward()
 
         # Compare gradients for all conversions
-        tol = 1e-4 if default_dtype == torch.float32 else 1e-8
+        tol = 1e-4 if default_dtype == torch.float32 else 1e-7
 
         def print_gradient_diff(name1, p1, name2, p2, conv_type):
             if p1.grad is not None and p1.grad.shape == p2.grad.shape:
@@ -152,7 +152,7 @@ def print_gradient_diff(name1, p1, name2, p2, conv_type):
                     print(
                         f"{conv_type} - Parameter {name1}/{name2}, Max error: {error.max()}"
                     )
-                    torch.testing.assert_close(p1.grad, p2.grad, atol=tol, rtol=1e-10)
+                    torch.testing.assert_close(p1.grad, p2.grad, atol=tol, rtol=tol)
 
         # E3nn to CuEq gradients
         for (name_e3nn, p_e3nn), (name_cueq, p_cueq) in zip(