update jaxtyped decorator (#44)

* update jaxtyped decorator

* add newline character to pacify flake

* add precomit hooks

* add missing return statement

* replace typing_extensions>=4.2.0 with typing_extensions

* pin version range for typing_extensions

* set min version of typing-extensions but not max

* bump version number to 0.2.8

MANIFEST.in

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

jaxtyping/__init__.py

@@ -20,14 +20,18 @@
 import typing
 
 
-if getattr(typing, "GENERATING_DOCUMENTATION", False):
-
-    class Array:
-        pass
-
-    Array.__module__ = "builtins"
-else:
+if typing.TYPE_CHECKING:
+    # type checkers don't know which branch below will be executed
     from jax.numpy import ndarray as Array
+else:
+    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,
@@ -63,4 +67,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.2.3"
+__version__ = "0.2.8"

jaxtyping/array_types.py

@@ -150,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,
@@ -383,10 +383,15 @@ class _MetaAbstractDtype(type):
                 elem = compile(elem, "<string>", "eval")
                 elem = _SymbolicDim(elem, broadcastable)
             dims.append(elem)
+        # In python 3.8, e.g., typing.Union lacks `__name__`.
+        try:
+            type_str = array_type.__name__
+        except AttributeError:
+            type_str = repr(array_type)
         if _array_name_format == "dtype_and_shape":
-            name = f"{cls.__name__}[{array_type.__name__}, '{dim_str}']"
+            name = f"{cls.__name__}[{type_str}, '{dim_str}']"
         elif _array_name_format == "array":
-            name = array_type.__name__
+            name = type_str
         else:
             raise ValueError(f"array_name_format {_array_name_format} not recognised")
         out = _MetaAbstractArray(
@@ -474,7 +479,10 @@ else:
 
         _Cls.__name__ = name
         _Cls.__qualname__ = name
-        _Cls.__module__ = "jaxtyping"
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            _Cls.__module__ = "builtins"
+        else:
+            _Cls.__module__ = "jaxtyping"
         return _Cls
 
     UInt8 = _make_dtype(_uint8, "UInt8")

jaxtyping/decorator.py

@@ -18,6 +18,7 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import functools as ft
+import inspect
 import threading
 
 
@@ -44,4 +45,9 @@ class _Jaxtyped:
 
 
 def jaxtyped(fn):
-    return ft.wraps(fn)(_Jaxtyped(fn))
+    if inspect.isclass(fn):  # allow decorators on class definitions
+        init = jaxtyped(fn.__init__)
+        fn.__init__ = init
+        return fn
+    else:
+        return ft.wraps(fn)(_Jaxtyped(fn))

jaxtyping/import_hook.py

@@ -65,7 +65,9 @@ def _call_with_frames_removed(f, *args, **kwargs):
 
 
 def _optimized_cache_from_source(path, debug_override=None):
-    return cache_from_source(path, debug_override, optimization="jaxtyping")
+    # Version 2: change the position of the `@jaxtyped` decorator, so need a
+    # different name to avoid hitting old __pycache__
+    return cache_from_source(path, debug_override, optimization="jaxtyping2")
 
 
 class _JaxtypingTransformer(ast.NodeVisitor):
@@ -99,10 +101,16 @@ class _JaxtypingTransformer(ast.NodeVisitor):
         has_annotated_args = any(arg for arg in node.args.args if arg.annotation)
         has_annotated_return = bool(node.returns)
         if has_annotated_args or has_annotated_return:
-            # Place at the start of the decorator list, in case a typechecking
-            # annotation has been manually applied; we need to be above that.
-            node.decorator_list.insert(
-                0,
+            # Place at the end of the decorator list, as otherwise we wrap e.g.
+            # `jax.custom_{jvp,vjp}` and lose the ability to `defjvp` etc.
+            #
+            # Note that the counter-argument here is that we'd like to place this
+            # at the start of the decorator list, in case a typechecking annotation
+            # has been manually applied, and we'd need to be above that. In this
+            # case we're just going to have to need to ask the user to remove their
+            # typechecking annotation (and let this decorator do it instead).
+            # It's more important we be compatible with normal JAX code.
+            node.decorator_list.append(
                 ast.Attribute(
                     ast.Name(id="jaxtyping", ctx=ast.Load()), "jaxtyped", ast.Load()
                 ),
@@ -230,8 +238,16 @@ def install_import_hook(
     - `typechecker`: the module and function of the typechecker you want to use, as a
         2-tuple of strings. For example `typechecker=("typeguard", "typechecked")` or
         `typechecker=("beartype", "beartype")`. You may pass `typechecker=None` if you
-        do not want to automatically decorate with a typechecker as well; e.g. if you
-        have a codebase that already has these decorators.
+        do not want to automatically decorate with a typechecker as well.
+
+    If the function already has any decorators on it, then both the `@jaxtyped` and the
+    typechecker decorators will go at the bottom of the decorator list, e.g.
+    ```python
+    @some_other_decorator
+    @jaxtyped
+    @beartype.beartype
+    def foo(...): ...
+    ```
 
     **Returns:**
 
@@ -243,8 +259,8 @@ def install_import_hook(
     ```python
     # entry_point.py
     from jaxtyped import install_import_hook
-    install_import_hook("main", ("beartype", "beartype"))
-    import main
+    with install_import_hook("main", ("beartype", "beartype"))
+        import main
     ...  # do whatever you're doing
 
     # main.py

jaxtyping/pytree_type.py

@@ -18,9 +18,11 @@
 # 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 jax.tree_util as jtu
 import typeguard
 
 
@@ -51,7 +53,10 @@ class _MetaPyTree(type):
     def __getitem__(cls, item):
         name = str(_FakePyTree[item])
         out = _MetaSubscriptPyTree(name, (), {"leaftype": item})
-        out.__module__ = "jaxtyping"
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            out.__module__ = "builtins"
+        else:
+            out.__module__ = "jaxtyping"
         return out
 
 
@@ -78,12 +83,22 @@ class _MetaSubscriptPyTree(type):
             else:
                 return True
 
-        leaves = jax.tree_leaves(obj, is_leaf=is_leaftype)
+        leaves = jtu.tree_leaves(obj, is_leaf=is_leaftype)
         return all(map(is_leaftype, leaves))
 
 
-PyTree = _MetaPyTree("PyTree", (), {})
-PyTree.__module__ = "jaxtyping"
+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

@@ -67,11 +67,14 @@ python_requires = "~=3.7"
 
 # We use typeguard internally (in a fairly minimal way), but it's not required that
 # end users make the same choice.
+# For typing_extensions, we choose versions that match
+# https://github.com/explosion/confection/blob/main/setup.cfg#L33 used in colab
+
 install_requires = [
     "jax>=0.3.4",
     "numpy>=1.20.0",
     "typeguard>=2.13.3",
-    "typing_extensions>=4.2.0",
+    "typing_extensions>=3.7.4.1",
 ]
 
 entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"])
@@ -94,4 +97,5 @@ setuptools.setup(
     install_requires=install_requires,
     entry_points=entry_points,
     packages=[name],
+    include_package_data=True,
 )

test/test_pytree.py

@@ -17,7 +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.
 
-from typing import Tuple, Union
+from typing import NamedTuple, Tuple, Union
 
 import equinox as eqx
 import jax
@@ -155,3 +155,33 @@ def test_pytree_tuple(typecheck):
         g([1, 1])
     with pytest.raises(ParamError):
         g([(1, 1), "hi"])
+
+
+def test_pytree_namedtuple(typecheck):
+    class CustomNamedTuple(NamedTuple):
+        x: Float[jnp.ndarray, "a b"]
+        y: Float[jnp.ndarray, "b c"]
+
+    class OtherCustomNamedTuple(NamedTuple):
+        x: Float[jnp.ndarray, "a b"]
+        y: Float[jnp.ndarray, "b c"]
+
+    @typecheck
+    def g(x: PyTree[CustomNamedTuple]):
+        ...
+
+    g(
+        CustomNamedTuple(
+            x=jax.random.normal(jax.random.PRNGKey(42), (3, 2)),
+            y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)),
+        )
+    )
+    with pytest.raises(ParamError):
+        g(object())
+    with pytest.raises(ParamError):
+        g(
+            OtherCustomNamedTuple(
+                x=jax.random.normal(jax.random.PRNGKey(42), (3, 2)),
+                y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)),
+            )
+        )

test/test_threading.py

@@ -25,7 +25,22 @@ from typeguard import typechecked
 from jaxtyping import Array, Float, jaxtyped
 
 
-def test_threading():
+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"]:
@@ -36,6 +51,16 @@ def test_threading():
         b = jnp.array([[2.0, 3.0]])
         add(a, b)
 
-    thread = threading.Thread(target=run)
+    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()