Now works with torch.compile? (#72)

jaxtyping/decorator.py

@@ -21,35 +21,20 @@ import dataclasses
 import functools as ft
 import inspect
 import threading
+import types
+import weakref
 
 
 storage = threading.local()
 
 
-class _Jaxtyped:
-    def __get__(self, instance, owner):
-        fn = self.__wrapped__
-        got = fn.__get__(instance, owner)
-        if fn is got:
-            return self
-        else:
-            return ft.wraps(got)(_Jaxtyped())
-
-    def __call__(self, *args, **kwargs):
-        try:
-            memo_stack = storage.memo_stack
-        except AttributeError:
-            memo_stack = storage.memo_stack = []
-        memo_stack.append(({}, {}, {}))
-        fn = self.__wrapped__
-        try:
-            return fn(*args, **kwargs)
-        finally:
-            memo_stack.pop()
+_jaxtyped_fns = weakref.WeakSet()
 
 
 def jaxtyped(fn):
-    if inspect.isclass(fn):  # allow decorators on class definitions
+    if type(fn) is types.FunctionType and fn in _jaxtyped_fns:
+        return fn
+    elif inspect.isclass(fn):  # allow decorators on class definitions
         if dataclasses.is_dataclass(fn):
             init = jaxtyped(fn.__init__)
             fn.__init__ = init
@@ -58,8 +43,44 @@ def jaxtyped(fn):
             raise ValueError(
                 "jaxtyped may only be added as a class decorator to dataclasses"
             )
+    # It'd be lovely if we could handle arbitrary descriptors, and not just the builtin
+    # ones. Unfortunately that means returning a class instance with a __get__ method,
+    # and that turns out to break loads of other things. See beartype issue #211 and
+    # jaxtyping issue #71.
+    elif isinstance(fn, classmethod):
+        return classmethod(jaxtyped(fn.__func__))
+    elif isinstance(fn, staticmethod):
+        return staticmethod(jaxtyped(fn.__func__))
+    elif isinstance(fn, property):
+        if fn.fget is None:
+            fget = None
+        else:
+            fget = jaxtyped(fn.fget)
+        if fn.fset is None:
+            fset = None
+        else:
+            fset = jaxtyped(fn.fset)
+        if fn.fdel is None:
+            fdel = None
+        else:
+            fdel = jaxtyped(fn.fdel)
+        return property(fget=fget, fset=fset, fdel=fdel)
     else:
-        return ft.wraps(fn)(_Jaxtyped())
+
+        @ft.wraps(fn)
+        def wrapped_fn(*args, **kwargs):
+            try:
+                memo_stack = storage.memo_stack
+            except AttributeError:
+                memo_stack = storage.memo_stack = []
+            memo_stack.append(({}, {}, {}))
+            try:
+                return fn(*args, **kwargs)
+            finally:
+                memo_stack.pop()
+
+        _jaxtyped_fns.add(wrapped_fn)
+        return wrapped_fn
 
 
 def _jaxtyped_typechecker(typechecker):
@@ -72,10 +93,8 @@ def _jaxtyped_typechecker(typechecker):
     def _wrapper(kls):
         assert inspect.isclass(kls)
         if dataclasses.is_dataclass(kls):
-            if type(kls.__init__) is not _Jaxtyped:
-                # Extra `if` check to work around beartype bug #211
-                init = jaxtyped(typechecker(kls.__init__))
-                kls.__init__ = init
+            init = jaxtyped(typechecker(kls.__init__))
+            kls.__init__ = init
         return kls
 
     return _wrapper

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