no global oeq import

nequix/model.py

@@ -11,22 +11,6 @@
 
 from nequix.layer_norm import RMSLayerNorm
 
-try:
-    import torch  # noqa: F401
-except ImportError:
-    # allow openequivariance to be imported without torch, but only if it is not
-    # installed; otherwise, the torch backend won't work if users want to use
-    # both torch and jax.
-    os.environ["OEQ_NOTORCH"] = "1"
-
-try:
-    import openequivariance as oeq
-    import openequivariance_extjax  # noqa: F401
-
-    OEQ_AVAILABLE = True
-except ImportError:
-    OEQ_AVAILABLE = False
-
 
 def bessel_basis(x: jax.Array, num_basis: int, r_max: float) -> jax.Array:
     prefactor = 2.0 / r_max
@@ -179,14 +163,26 @@ def __init__(
         self.tp_irreps = tp_irreps
 
         if kernel:
-            instructions = [instructions[i] for i in inv]
-            if not OEQ_AVAILABLE:
+            try:
+                import torch  # noqa: F401
+            except ImportError:
+                # allow openequivariance to be imported without torch, but only if it is not
+                # installed; otherwise, the torch backend won't work if users want to use
+                # both torch and jax.
+                os.environ["OEQ_NOTORCH"] = "1"
+
+            try:
+                import openequivariance as oeq
+                import openequivariance_extjax  # noqa: F401
+            except ImportError:
                 raise ImportError(
                     "OpenEquivariance with JAX support is required for kernel=True. "
                     "Install both packages:\n"
                     "  uv pip install 'openequivariance[jax]'\n"
                     "  uv pip install 'openequivariance_extjax' --no-build-isolation"
                 )
+
+            instructions = [instructions[i] for i in inv]
             problem = oeq.TPProblem(
                 str(input_irreps),
                 str(sh_irreps),

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nequix"
-version = "0.4.4"
+version = "0.4.5"
 description = "Nequix source code"
 readme = "README.md"
 authors = [

tests/test_calculator.py

@@ -4,7 +4,14 @@
 import ase.build
 
 from nequix.calculator import NequixCalculator
-from nequix.model import OEQ_AVAILABLE
+
+try:
+    import openequivariance  # noqa: F401
+    import openequivariance_extjax  # noqa: F401
+
+    OEQ_AVAILABLE = True
+except (ImportError, AssertionError):
+    OEQ_AVAILABLE = False
 
 skip_no_oeq = pytest.mark.skipif(not OEQ_AVAILABLE, reason="OpenEquivariance not installed")
 

tests/test_jax_md.py

@@ -8,7 +8,7 @@
 
 from nequix.calculator import NequixCalculator
 from nequix.data import atomic_numbers_to_indices
-from nequix.model import OEQ_AVAILABLE
+from tests.test_calculator import OEQ_AVAILABLE
 
 try:
     from jax_md import quantity, space