From 9c900dc4b91b24e0d688ed42808a15668cc07dea Mon Sep 17 00:00:00 2001 From: Owen Lockwood <42878312+lockwo@users.noreply.github.com> Date: Mon, 25 May 2026 17:33:51 -0700 Subject: [PATCH 1/3] strict fix --- jumpax/_aggregators/base.py | 2 +- jumpax/_custom_meta.py | 29 ----------------------------- jumpax/_jumps.py | 2 +- pyproject.toml | 2 +- tests/test_strict_meta.py | 2 +- 5 files changed, 4 insertions(+), 33 deletions(-) delete mode 100644 jumpax/_custom_meta.py diff --git a/jumpax/_aggregators/base.py b/jumpax/_aggregators/base.py index 955439b..7ce4a50 100644 --- a/jumpax/_aggregators/base.py +++ b/jumpax/_aggregators/base.py @@ -4,7 +4,7 @@ import equinox as eqx from jaxtyping import Array, Bool, Int, Key -from .._custom_meta import AbstractStrictModule +from ihoop.eqx import AbstractStrictModule from .._custom_types import Args, Info, JumpState, RealScalarLike, SolverState, U _SolverState = TypeVar("_SolverState", bound=SolverState) diff --git a/jumpax/_custom_meta.py b/jumpax/_custom_meta.py deleted file mode 100644 index 8bfb547..0000000 --- a/jumpax/_custom_meta.py +++ /dev/null @@ -1,29 +0,0 @@ -# type: ignore -from abc import abstractmethod - -import equinox as eqx -import ihoop - - -class _StrictEqxMeta(ihoop.strict._StrictMeta, eqx._module._module._ModuleMeta): - def __new__(mcs, name, bases, namespace, **kwargs): - if name == "AbstractStrictModule": - - @abstractmethod - def _strict_base_(self): - raise NotImplementedError - - namespace["_strict_base_"] = _strict_base_ - elif not name.startswith("Abstract") and not name.startswith("_Abstract"): - if "_strict_base_" not in namespace: - - def _strict_base_(self): - pass - - namespace["_strict_base_"] = _strict_base_ - return super().__new__(mcs, name, bases, namespace, **kwargs) - - -class AbstractStrictModule(eqx.Module, ihoop.Strict, metaclass=_StrictEqxMeta): - def __init_subclass__(cls, *, strict: bool = False, **kwargs): - super().__init_subclass__(**kwargs) diff --git a/jumpax/_jumps.py b/jumpax/_jumps.py index a0ae3c1..3e013e1 100644 --- a/jumpax/_jumps.py +++ b/jumpax/_jumps.py @@ -7,7 +7,7 @@ from jax.scipy.special import gammaln from jaxtyping import Array, Float, Int, Key -from ._custom_meta import AbstractStrictModule +from ihoop.eqx import AbstractStrictModule from ._custom_types import Args, JumpState, Rate, RealScalarLike, U _Rate = TypeVar("_Rate", bound=Rate) diff --git a/pyproject.toml b/pyproject.toml index e58f58f..6e00b78 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/tests/test_strict_meta.py b/tests/test_strict_meta.py index 90a204a..5e2a972 100644 --- a/tests/test_strict_meta.py +++ b/tests/test_strict_meta.py @@ -3,7 +3,7 @@ from abc import abstractmethod from unittest import TestCase -from jumpax._custom_meta import AbstractStrictModule +from ihoop.eqx import AbstractStrictModule class StrictMetaTest(TestCase): From da070cf35d59d09992d1f8cd00027aaa79a893c6 Mon Sep 17 00:00:00 2001 From: Owen Lockwood <42878312+lockwo@users.noreply.github.com> Date: Mon, 25 May 2026 17:37:32 -0700 Subject: [PATCH 2/3] format --- jumpax/_aggregators/base.py | 2 +- jumpax/_jumps.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/jumpax/_aggregators/base.py b/jumpax/_aggregators/base.py index 7ce4a50..1581a58 100644 --- a/jumpax/_aggregators/base.py +++ b/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 ihoop.eqx import AbstractStrictModule from .._custom_types import Args, Info, JumpState, RealScalarLike, SolverState, U _SolverState = TypeVar("_SolverState", bound=SolverState) diff --git a/jumpax/_jumps.py b/jumpax/_jumps.py index 3e013e1..d0b2867 100644 --- a/jumpax/_jumps.py +++ b/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 ihoop.eqx import AbstractStrictModule from ._custom_types import Args, JumpState, Rate, RealScalarLike, U _Rate = TypeVar("_Rate", bound=Rate) From da35d5566d4d1071f904cb3281a6af5de2fb8e6e Mon Sep 17 00:00:00 2001 From: Owen Lockwood <42878312+lockwo@users.noreply.github.com> Date: Mon, 25 May 2026 17:40:24 -0700 Subject: [PATCH 3/3] a --- tests/test_strict_meta.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_strict_meta.py b/tests/test_strict_meta.py index 5e2a972..5bd0edd 100644 --- a/tests/test_strict_meta.py +++ b/tests/test_strict_meta.py @@ -10,24 +10,24 @@ 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