strict fix

jumpax/_aggregators/base.py

@@ -2,9 +2,9 @@
 from typing import Generic, TypeVar
 
 import equinox as eqx
+from ihoop.eqx import AbstractStrictModule
 from jaxtyping import Array, Bool, Int, Key
 
-from .._custom_meta import AbstractStrictModule
 from .._custom_types import Args, Info, JumpState, RealScalarLike, SolverState, U
 
 _SolverState = TypeVar("_SolverState", bound=SolverState)

jumpax/_jumps.py

@@ -3,11 +3,11 @@
 from typing import Generic, TypeAlias, TypeVar
 
 import jax
+from ihoop.eqx import AbstractStrictModule
 from jax import numpy as jnp
 from jax.scipy.special import gammaln
 from jaxtyping import Array, Float, Int, Key
 
-from ._custom_meta import AbstractStrictModule
 from ._custom_types import Args, JumpState, Rate, RealScalarLike, U
 
 _Rate = TypeVar("_Rate", bound=Rate)

pyproject.toml

@@ -24,7 +24,7 @@ urls = { repository = "https://github.com/lockwo/jumpax" }
 
 dependencies = [
     "equinox>=0.11.2",
-    "ihoop>=0.1.1",
+    "ihoop[equinox]>=0.1.3",
     "diffrax>=0.6.0",
 ]
 

tests/test_strict_meta.py

@@ -3,31 +3,31 @@
 from abc import abstractmethod
 from unittest import TestCase
 
-from jumpax._custom_meta import AbstractStrictModule
+from ihoop.eqx import AbstractStrictModule
 
 
 class StrictMetaTest(TestCase):
     def test_concrete_class_is_final(self):
         """Concrete strict classes cannot be subclassed."""
 
-        class AbstractFoo(AbstractStrictModule, strict=True):
+        class AbstractFoo(AbstractStrictModule):
             @abstractmethod
             def bar(self):
                 raise NotImplementedError
 
-        class Foo(AbstractFoo, strict=True):
+        class Foo(AbstractFoo):
             def bar(self):
                 return 42
 
         with self.assertRaises(TypeError, msg="Concrete classes must be final"):
 
-            class SubFoo(Foo, strict=True):
+            class SubFoo(Foo):
                 pass
 
     def test_abstract_class_cannot_be_instantiated(self):
         """Abstract strict classes cannot be instantiated."""
 
-        class AbstractFoo(AbstractStrictModule, strict=True):
+        class AbstractFoo(AbstractStrictModule):
             @abstractmethod
             def bar(self):
                 raise NotImplementedError