Fixed jaxtyped breaking descriptors. Fixed long module names. (#25)

jaxtyping/__init__.py

@@ -53,4 +53,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"

jaxtyping/array_types.py

@@ -388,7 +388,7 @@ class _MetaAbstractDtype(type):
             name = "Array"
         else:
             raise ValueError(f"array_name_format {_array_name_format} not recognised")
-        return _MetaAbstractArray(
+        out = _MetaAbstractArray(
             name,
             (AbstractArray,),
             dict(
@@ -398,6 +398,8 @@ class _MetaAbstractDtype(type):
                 index_variadic=index_variadic,
             ),
         )
+        out.__module__ = "jaxtyping"
+        return out
 
 
 class AbstractDtype(metaclass=_MetaAbstractDtype):
@@ -468,6 +470,7 @@ else:
 
         _Cls.__name__ = name
         _Cls.__qualname__ = name
+        _Cls.__module__ = "jaxtyping"
         return _Cls
 
     UInt8 = _make_dtype(_uint8, "UInt8")

jaxtyping/decorator.py

@@ -22,16 +22,26 @@ import threading
 
 
 storage = threading.local()
-storage.memo_stack = []
+
+
+class _Jaxtyped:
+    def __init__(self, fn):
+        self.fn = fn
+
+    def __get__(self, instance, owner):
+        return ft.wraps(self.fn)(_Jaxtyped(self.fn.__get__(instance, owner)))
+
+    def __call__(self, *args, **kwargs):
+        try:
+            memo_stack = storage.memo_stack
+        except AttributeError:
+            memo_stack = storage.memo_stack = []
+        memo_stack.append(({}, {}, {}))
+        try:
+            return self.fn(*args, **kwargs)
+        finally:
+            memo_stack.pop()
 
 
 def jaxtyped(fn):
-    @ft.wraps(fn)
-    def wrapper(*args, **kwargs):
-        storage.memo_stack.append(({}, {}, {}))
-        try:
-            return fn(*args, **kwargs)
-        finally:
-            storage.memo_stack.pop()
-
-    return wrapper
+    return ft.wraps(fn)(_Jaxtyped(fn))

jaxtyping/pytree_type.py

@@ -33,6 +33,7 @@ class _FakePyTree(Generic[_T]):
 
 _FakePyTree.__name__ = "PyTree"
 _FakePyTree.__qualname__ = "PyTree"
+_FakePyTree.__module__ = "builtins"
 # Can't do type("PyTree", (Generic[_T],), {}) because dynamic subclassing of typeforms
 # isn't allowed.
 # Can't do types.new_class("PyTree", (Generic[_T],), {}) because that has __module__
@@ -49,7 +50,9 @@ class _MetaPyTree(type):
     @ft.lru_cache(maxsize=None)
     def __getitem__(cls, item):
         name = str(_FakePyTree[item])
-        return _MetaSubscriptPyTree(name, (), {"leaftype": item})
+        out = _MetaSubscriptPyTree(name, (), {"leaftype": item})
+        out.__module__ = "jaxtyping"
+        return out
 
 
 class _MetaSubscriptPyTree(type):
@@ -80,6 +83,7 @@ class _MetaSubscriptPyTree(type):
 
 
 PyTree = _MetaPyTree("PyTree", (), {})
+PyTree.__module__ = "jaxtyping"
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
 # instancecheck for PyTree[foo], but we subclassing
 # `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.

test/test_decorator.py

@@ -0,0 +1,13 @@
+from jaxtyping import jaxtyped
+
+
+class M:
+    @jaxtyped
+    @classmethod
+    def f(cls):
+        return 3
+
+
+# Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod
+def test_decorator():
+    assert M.f() == 3

test/test_threading.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2022 Google LLC
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to
+# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+# the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+# 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 threading
+
+import jax.numpy as jnp
+from typeguard import typechecked
+
+from jaxtyping import Array, Float, jaxtyped
+
+
+def test_threading():
+    @jaxtyped
+    @typechecked
+    def add(x: Float[Array, "a b"], y: Float[Array, "a b"]) -> Float[Array, "a b"]:
+        return x + y
+
+    def run():
+        a = jnp.array([[1.0, 2.0]])
+        b = jnp.array([[2.0, 3.0]])
+        add(a, b)
+
+    thread = threading.Thread(target=run)
+    thread.start()
+    thread.join()