Now works with torch.compile? (#72)

test/import_hook_tester_typeguard.py

@@ -17,9 +17,7 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-import dataclasses
 
-import equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -37,26 +35,29 @@ with pytest.raises(ParamError):
     g(jnp.array(1))
 
 
-class M(eqx.Module):
-    foo: int
-    bar: Float32[jnp.ndarray, " a"]
-
-
-M(1, jnp.array([1.0]))
-with pytest.raises(ParamError):
-    M(1.0, jnp.array([1.0]))
-with pytest.raises(ParamError):
-    M(1, jnp.array(1.0))
-
-
-@dataclasses.dataclass
-class D:
-    foo: int
-    bar: Float32[jnp.ndarray, " a"]
-
-
-D(1, jnp.array([1.0]))
-with pytest.raises(ParamError):
-    D(1.0, jnp.array([1.0]))
-with pytest.raises(ParamError):
-    D(1, jnp.array(1.0))
+# Typeguard 3.0 no longer supports this.
+#
+# class M(eqx.Module):
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# M(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1, jnp.array(1.0))
+#
+#
+#
+# @dataclasses.dataclass
+# class D:
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# D(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1, jnp.array(1.0))

test/import_hook_tester_typeguard_old.py

@@ -17,9 +17,7 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-import dataclasses
 
-import equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -37,26 +35,29 @@ with pytest.raises(ParamError):
     g(jnp.array(1))
 
 
-class M(eqx.Module):
-    foo: int
-    bar: Float32[jnp.ndarray, " a"]
-
-
-M(1, jnp.array([1.0]))
-with pytest.raises(ParamError):
-    M(1.0, jnp.array([1.0]))
-with pytest.raises(ParamError):
-    M(1, jnp.array(1.0))
-
-
-@dataclasses.dataclass
-class D:
-    foo: int
-    bar: Float32[jnp.ndarray, " a"]
-
-
-D(1, jnp.array([1.0]))
-with pytest.raises(ParamError):
-    D(1.0, jnp.array([1.0]))
-with pytest.raises(ParamError):
-    D(1, jnp.array(1.0))
+# Typeguard 3.0 no longer supports this.
+#
+# class M(eqx.Module):
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# M(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1, jnp.array(1.0))
+#
+#
+#
+# @dataclasses.dataclass
+# class D:
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# D(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1, jnp.array(1.0))

test/test_decorator.py

@@ -4,43 +4,74 @@ from jaxtyping import jaxtyped
 
 
 class M(metaclass=abc.ABCMeta):
+    @jaxtyped
+    def f(self):
+        ...
+
     @jaxtyped
     @classmethod
-    def f1(cls):
+    def g1(cls):
         return 3
 
     @classmethod
     @jaxtyped
-    def f2(cls):
+    def g2(cls):
+        return 4
+
+    @jaxtyped
+    @staticmethod
+    def h1():
+        return 3
+
+    @staticmethod
+    @jaxtyped
+    def h2():
         return 4
 
     @jaxtyped
     @abc.abstractmethod
-    def g1(self):
+    def i1(self):
         ...
 
     @abc.abstractmethod
     @jaxtyped
-    def g2(self):
+    def i2(self):
         ...
 
+
+class N:
     @jaxtyped
-    def h(self):
-        ...
+    @property
+    def j1(self):
+        return 3
+
+    @property
+    @jaxtyped
+    def j2(self):
+        return 4
+
+
+def test_identity():
+    assert M.f is M.f
 
 
-# Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod
 def test_classmethod():
-    assert M.f1() == 3
-    assert M.f2() == 4
+    assert M.g1() == 3
+    assert M.g2() == 4
+
+
+def test_staticmethod():
+    assert M.h1() == 3
+    assert M.h2() == 4
 
 
 # Check that the @jaxtyped decorator doesn't blat the __isabstractmethod__ of
 # @abstractmethod
 def test_abstractmethod():
-    assert M.g1.__isabstractmethod__
-    assert M.g2.__isabstractmethod__
+    assert M.i1.__isabstractmethod__
+    assert M.i2.__isabstractmethod__
 
 
-def test_identity():
-    assert M.h is M.h
+def test_property():
+    assert N().j1 == 3
+    assert N().j2 == 4

test/test_import_hook.py

@@ -78,7 +78,7 @@ def test_import_hook_beartype_full():
 
 def test_import_hook_transitive():
     hook = install_import_hook(
-        "test.import_hook_tester_transitive", "typeguard.typechecked"
+        "test.import_hook_tester_transitive", "beartype.beartype"
     )
     with hook:
         from . import import_hook_tester_transitive  # noqa: F401