Fixed py.type not being packaged (#30)

* Edge-case doc fixes for parameterising types with PyTrees or AbstractDtypes

* Fixes for threading

* version bump

MANIFEST.in

@@ -1,2 +1,2 @@
-include LICENSE
-prune tests
+include jaxtyping/py.typed
+prune test

jaxtyping/__init__.py

@@ -17,7 +17,17 @@
 # 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.
 
-from jax.numpy import ndarray as Array
+import typing
+
+
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+    class Array:
+        pass
+
+    Array.__module__ = "builtins"
+else:
+    from jax.numpy import ndarray as Array
 
 from .array_types import (
     AbstractArray,
@@ -53,4 +63,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.2.0"
+__version__ = "0.2.5"

jaxtyping/array_types.py

@@ -19,6 +19,7 @@
 
 import enum
 import functools as ft
+import typing
 from typing import Any, Dict, List, NoReturn, Optional, Tuple, TYPE_CHECKING, Union
 from typing_extensions import Literal
 
@@ -149,25 +150,25 @@ class _MetaAbstractArray(type):
         if cls.dtypes is not _any_dtype and dtype not in cls.dtypes:
             return False
 
-        if len(storage.memo_stack) == 0:
+        no_temp_memo = hasattr(storage, "memo_stack") and len(storage.memo_stack) != 0
+
+        if no_temp_memo:
+            single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]
+            # Make a copy so we don't mutate the original memo during the shape check.
+            single_memo = single_memo.copy()
+            variadic_memo = variadic_memo.copy()
+            variadic_broadcast_memo = variadic_broadcast_memo.copy()
+        else:
             # `isinstance` happening outside any @jaxtyped decorators, e.g. at the
             # global scope. In this case just create a temporary memo, since we're not
             # going to be comparing against any stored values anyway.
             single_memo = {}
             variadic_memo = {}
             variadic_broadcast_memo = {}
-            temp_memo = True
-        else:
-            single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]
-            # Make a copy so we don't mutate the original memo during the shape check.
-            single_memo = single_memo.copy()
-            variadic_memo = variadic_memo.copy()
-            variadic_broadcast_memo = variadic_broadcast_memo.copy()
-            temp_memo = False
 
         if cls._check_shape(obj, single_memo, variadic_memo, variadic_broadcast_memo):
             # We update the memo every time we successfully pass a shape check
-            if not temp_memo:
+            if no_temp_memo:
                 storage.memo_stack[-1] = (
                     single_memo,
                     variadic_memo,
@@ -385,10 +386,10 @@ class _MetaAbstractDtype(type):
         if _array_name_format == "dtype_and_shape":
             name = f"{cls.__name__}[{array_type.__name__}, '{dim_str}']"
         elif _array_name_format == "array":
-            name = "Array"
+            name = array_type.__name__
         else:
             raise ValueError(f"array_name_format {_array_name_format} not recognised")
-        return _MetaAbstractArray(
+        out = _MetaAbstractArray(
             name,
             (AbstractArray,),
             dict(
@@ -398,6 +399,11 @@ class _MetaAbstractDtype(type):
                 index_variadic=index_variadic,
             ),
         )
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            out.__module__ = "builtins"
+        else:
+            out.__module__ = "jaxtyping"
+        return out
 
 
 class AbstractDtype(metaclass=_MetaAbstractDtype):
@@ -446,7 +452,7 @@ if TYPE_CHECKING:
     from typing_extensions import Annotated as UInt32
     from typing_extensions import Annotated as UInt64
 else:
-    _bool = "bool"
+    _bool = "bool_"
     _uint8 = "uint8"
     _uint16 = "uint16"
     _uint32 = "uint32"
@@ -468,6 +474,10 @@ else:
 
         _Cls.__name__ = name
         _Cls.__qualname__ = name
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            _Cls.__module__ = "builtins"
+        else:
+            _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/py.typed

@@ -0,0 +1 @@
+

jaxtyping/pytree_type.py

@@ -18,7 +18,9 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import functools as ft
-from typing import Generic, TypeVar
+import typing
+from typing import Generic, TYPE_CHECKING, TypeVar
+from typing_extensions import Protocol
 
 import jax
 import typeguard
@@ -33,6 +35,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 +52,12 @@ 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})
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            out.__module__ = "builtins"
+        else:
+            out.__module__ = "jaxtyping"
+        return out
 
 
 class _MetaSubscriptPyTree(type):
@@ -79,7 +87,18 @@ class _MetaSubscriptPyTree(type):
         return all(map(is_leaftype, leaves))
 
 
-PyTree = _MetaPyTree("PyTree", (), {})
+if TYPE_CHECKING:
+    # Work around pytype bug #1288
+    # pytype: skip-file
+    class PyTree(Protocol[_T]):
+        pass
+
+else:
+    PyTree = _MetaPyTree("PyTree", (), {})
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        PyTree.__module__ = "builtins"
+    else:
+        PyTree.__module__ = "jaxtyping"
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
-# instancecheck for PyTree[foo], but we subclassing
+# instancecheck for PyTree[foo], but subclassing
 # `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.

setup.py

@@ -94,4 +94,5 @@ setuptools.setup(
     install_requires=install_requires,
     entry_points=entry_points,
     packages=[name],
+    include_package_data=True,
 )

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,66 @@
+# 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
+
+
+class _ErrorableThread(threading.Thread):
+    def run(self):
+        try:
+            super().run()
+        except Exception as e:
+            self.exc = e
+        finally:
+            del self._target, self._args, self._kwargs
+
+    def join(self, timeout=None):
+        super().join(timeout)
+        if hasattr(self, "exc"):
+            raise self.exc
+
+
+def test_threading_jaxtyped():
+    @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 = _ErrorableThread(target=run)
+    thread.start()
+    thread.join()
+
+
+def test_threading_nojaxtyped():
+    def run():
+        a = jnp.array([[1.0, 2.0]])
+        assert isinstance(a, Float[Array, "..."])
+
+    thread = _ErrorableThread(target=run)
+    thread.start()
+    thread.join()