Hopefully fixed PyTree raising spurious errors. Bit mysterious that this worked before, really. I've tested this fix as best I can against the various static type checkers, but these are weird and varied enough that this might not be a perfect fix. If you see this and have issues, let me know. (#54)

* Explicitly export names to make pyright happy

* Bump jax and jaxtyping versions

* Add note on `jaxtyping` names

* Remove __all__ from `array_types.py`

* Appease flake8

* Reduce import redundancy

* Fix capitalization

.isort.cfg

@@ -2,5 +2,6 @@
 force_alphabetical_sort_within_sections=true
 lines_after_imports=2
 profile=black
+combine_as_imports=True
 treat_comments_as_code=true
 extra_standard_library=typing_extensions

API.md

@@ -168,6 +168,8 @@ The import hook can be applied to multiple packages via
 install_import_hook(["foo", "bar.baz"], ...)
 ```
 
+The import hook will automatically decorate all functions, and the `__init__` method of dataclasses.
+
 **Example: writing an end-user script**
 
 ```python

MANIFEST.in

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

README.md

@@ -28,8 +28,6 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 pip install jaxtyping
 ```
 
-Requires JAX 0.3.4+.
-
 Also install your favourite runtime type-checking package. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments).
 
 ## Documentation
@@ -46,6 +44,8 @@ Neural networks: [Equinox](https://github.com/patrick-kidger/equinox).
 
 Numerical differential equation solvers: [Diffrax](https://github.com/patrick-kidger/diffrax).
 
+Computer vision models: [Eqxvision](https://github.com/paganpasta/eqxvision).
+
 SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2jax](https://github.com/google/sympy2jax).
 
 ### Acknowledgements

jaxtyping/__init__.py

@@ -17,40 +17,89 @@
 # 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
+import typing_extensions
+
+
+try:
+    import jax
+except ImportError:
+    has_jax = False
+else:
+    has_jax = True
+    del jax
+
+
+# Type checkers don't know which branch below will be executed.
+if typing.TYPE_CHECKING:
+    # For imports, we need to explicitly `import X as X` in order for Pyright to see
+    # them as public. See discussion at https://github.com/microsoft/pyright/issues/2277
+    from jax import Array as Array
+elif has_jax:
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+        class Array:
+            pass
+
+        Array.__module__ = "builtins"
+    else:
+        from jax import Array as Array
 
 from .array_types import (
-    AbstractArray,
-    AbstractDtype,
-    BFloat16,
-    Bool,
-    Complex,
-    Complex64,
-    Complex128,
-    Float,
-    Float16,
-    Float32,
-    Float64,
-    get_array_name_format,
-    Inexact,
-    Int,
-    Int8,
-    Int16,
-    Int32,
-    Int64,
-    Integer,
-    Num,
-    set_array_name_format,
-    Shaped,
-    UInt,
-    UInt8,
-    UInt16,
-    UInt32,
-    UInt64,
+    AbstractArray as AbstractArray,
+    AbstractDtype as AbstractDtype,
+    BFloat16 as BFloat16,
+    Bool as Bool,
+    Complex as Complex,
+    Complex64 as Complex64,
+    Complex128 as Complex128,
+    Float as Float,
+    Float16 as Float16,
+    Float32 as Float32,
+    Float64 as Float64,
+    get_array_name_format as get_array_name_format,
+    Inexact as Inexact,
+    Int as Int,
+    Int8 as Int8,
+    Int16 as Int16,
+    Int32 as Int32,
+    Int64 as Int64,
+    Integer as Integer,
+    Num as Num,
+    set_array_name_format as set_array_name_format,
+    Shaped as Shaped,
+    UInt as UInt,
+    UInt8 as UInt8,
+    UInt16 as UInt16,
+    UInt32 as UInt32,
+    UInt64 as UInt64,
 )
-from .decorator import jaxtyped
-from .import_hook import install_import_hook
-from .pytree_type import PyTree
+from .decorator import jaxtyped as jaxtyped
+from .import_hook import install_import_hook as install_import_hook
 
 
-__version__ = "0.2.1"
+if typing.TYPE_CHECKING:
+    # Set up to deliberately confuse a static type checker.
+    PyTree = getattr(typing, "foo" + "bar")
+    # What's going on with this madness?
+    #
+    # At static-type-checking-time, we want `PyTree` to be a type for which both
+    # `PyTree` and `PyTree[Foo]` are equivalent to `Any`.
+    # (The intention is that `PyTree` be a runtime-only type; there's no real way to
+    # do more with static type checkers.)
+    #
+    # Unfortunately, this isn't possible: `Any` isn't subscriptable. And there's no
+    # equivalent way we can fake this using typing annotations. (In some sense the
+    # closest thing would be a `Protocol[T]` with no methods, but that's actually the
+    # opposite of what we want: that ends up allowing nothing at all.)
+    #
+    # The good news for us is that static type checkers have an internal escape hatch.
+    # If they can't figure out what a type is, then they just give up and allow
+    # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking
+    # annotation, which static type checkers aren't smart enough to resolve.
+elif has_jax:
+    from .pytree_type import PyTree
+
+del has_jax
+
+__version__ = "0.2.11"

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,
@@ -382,13 +383,18 @@ 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"
+            name = type_str
         else:
             raise ValueError(f"array_name_format {_array_name_format} not recognised")
-        return _MetaAbstractArray(
+        out = _MetaAbstractArray(
             name,
             (AbstractArray,),
             dict(
@@ -398,6 +404,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):
@@ -422,29 +433,31 @@ if TYPE_CHECKING:
     # Note that `from typing_extensions import Annotated; ... = Annotated`
     # does not work with static type checkers. `Annotated` is a typeform rather
     # than a type, meaning it cannot be assigned.
-    from typing_extensions import Annotated as BFloat16
-    from typing_extensions import Annotated as Bool
-    from typing_extensions import Annotated as Complex
-    from typing_extensions import Annotated as Complex64
-    from typing_extensions import Annotated as Complex128
-    from typing_extensions import Annotated as Float
-    from typing_extensions import Annotated as Float16
-    from typing_extensions import Annotated as Float32
-    from typing_extensions import Annotated as Float64
-    from typing_extensions import Annotated as Inexact
-    from typing_extensions import Annotated as Int
-    from typing_extensions import Annotated as Int8
-    from typing_extensions import Annotated as Int16
-    from typing_extensions import Annotated as Int32
-    from typing_extensions import Annotated as Int64
-    from typing_extensions import Annotated as Integer
-    from typing_extensions import Annotated as Num
-    from typing_extensions import Annotated as Shaped
-    from typing_extensions import Annotated as UInt
-    from typing_extensions import Annotated as UInt8
-    from typing_extensions import Annotated as UInt16
-    from typing_extensions import Annotated as UInt32
-    from typing_extensions import Annotated as UInt64
+    from typing_extensions import (
+        Annotated as BFloat16,
+        Annotated as Bool,
+        Annotated as Complex,
+        Annotated as Complex64,
+        Annotated as Complex128,
+        Annotated as Float,
+        Annotated as Float16,
+        Annotated as Float32,
+        Annotated as Float64,
+        Annotated as Inexact,
+        Annotated as Int,
+        Annotated as Int8,
+        Annotated as Int16,
+        Annotated as Int32,
+        Annotated as Int64,
+        Annotated as Integer,
+        Annotated as Num,
+        Annotated as Shaped,
+        Annotated as UInt,
+        Annotated as UInt8,
+        Annotated as UInt16,
+        Annotated as UInt32,
+        Annotated as UInt64,
+    )
 else:
     _bool = "bool_"
     _uint8 = "uint8"
@@ -468,6 +481,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

@@ -17,21 +17,60 @@
 # 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 functools as ft
+import inspect
 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()
+    if inspect.isclass(fn):  # allow decorators on class definitions
+        if dataclasses.is_dataclass(fn):
+            init = jaxtyped(fn.__init__)
+            fn.__init__ = init
+            return fn
+        else:
+            raise ValueError(
+                "jaxtyped may only be added as a class decorator to dataclasses"
+            )
+    else:
+        return ft.wraps(fn)(_Jaxtyped(fn))
 
-    return wrapper
+
+def _jaxtyped_typechecker(typechecker):
+    # typechecker is expected to probably be either `typeguard.typechecked`, or
+    # `beartype.beartype`, or `None`.
+
+    if typechecker is None:
+        typechecker = lambda x: x
+
+    def _wrapper(kls):
+        assert inspect.isclass(kls)
+        if dataclasses.is_dataclass(kls):
+            init = jaxtyped(typechecker(kls.__init__))
+            kls.__init__ = init
+        return kls
+
+    return _wrapper

jaxtyping/import_hook.py

@@ -65,7 +65,17 @@ 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__
+    # Version 3: now also annotating classes.
+    return cache_from_source(path, debug_override, optimization="jaxtyping3")
+
+
+def _dot_lookup(*elements):
+    out = ast.Name(id=elements[0], ctx=ast.Load())
+    for element in elements[1:]:
+        out = ast.Attribute(out, element, ctx=ast.Load())
+    return out
 
 
 class _JaxtypingTransformer(ast.NodeVisitor):
@@ -95,31 +105,37 @@ class _JaxtypingTransformer(ast.NodeVisitor):
         self._parents.pop()
         return node
 
+    def visit_ClassDef(self, node: ast.ClassDef):
+        func = _dot_lookup("jaxtyping", "decorator", "_jaxtyped_typechecker")
+        if self._typechecker is None:
+            args = [ast.Constant(None)]
+        else:
+            args = [_dot_lookup(*self._typechecker)]
+        node.decorator_list.append(ast.Call(func, args, keywords=[]))
+        self._parents.append(node)
+        self.generic_visit(node)
+        self._parents.pop()
+        return node
+
     def visit_FunctionDef(self, node: ast.FunctionDef):
         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,
-                ast.Attribute(
-                    ast.Name(id="jaxtyping", ctx=ast.Load()), "jaxtyped", ast.Load()
-                ),
-            )
+            # 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(_dot_lookup("jaxtyping", "jaxtyped"))
             if self._typechecker is not None:
                 # Place at the end of the decorator list, as decorators
                 # frequently remove annotations from functions and we'd like to
                 # use those annotations.
-                typechecker_module, typechecker_function = self._typechecker
-                node.decorator_list.append(
-                    ast.Attribute(
-                        ast.Name(id=typechecker_module, ctx=ast.Load()),
-                        typechecker_function,
-                        ast.Load(),
-                    )
-                )
-
+                node.decorator_list.append(_dot_lookup(*self._typechecker))
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
@@ -230,8 +246,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 +267,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,10 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import functools as ft
+import typing
 from typing import Generic, TypeVar
 
-import jax
+import jax.tree_util as jtu
 import typeguard
 
 
@@ -33,6 +34,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 +51,20 @@ 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
+
+
+try:
+    # new typeguard
+    _TypeCheckError = (TypeError, typeguard.TypeCheckError)
+except AttributeError:
+    # old typeguard
+    _TypeCheckError = TypeError
 
 
 class _MetaSubscriptPyTree(type):
@@ -70,16 +85,20 @@ class _MetaSubscriptPyTree(type):
         def is_leaftype(x):
             try:
                 accepts_leaftype(x)
-            except TypeError:
+            except _TypeCheckError:
                 return False
             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", (), {})
 # 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.
+PyTree = _MetaPyTree("PyTree", (), {})
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+    PyTree.__module__ = "builtins"
+else:
+    PyTree.__module__ = "jaxtyping"

setup.py

@@ -67,11 +67,13 @@ 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 +96,5 @@ setuptools.setup(
     install_requires=install_requires,
     entry_points=entry_points,
     packages=[name],
+    include_package_data=True,
 )

test/helpers.py

@@ -18,16 +18,31 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import equinox as eqx
+import typeguard
 
 
+ParamError = []
+ReturnError = []
+ParamError.append(TypeError)  # old typeguard
+ReturnError.append(TypeError)  # old typeguard
+
+try:
+    # new typeguard
+    ParamError.append(typeguard.TypeCheckError)
+    ReturnError.append(typeguard.TypeCheckError)
+except AttributeError:
+    pass
+
 try:
     import beartype
 except ImportError:
-    ParamError = TypeError
-    ReturnError = TypeError
+    pass
 else:
-    ParamError = (TypeError, beartype.roar.BeartypeCallHintParamViolation)
-    ReturnError = (TypeError, beartype.roar.BeartypeCallHintReturnViolation)
+    ParamError.append(beartype.roar.BeartypeCallHintParamViolation)
+    ReturnError.append(beartype.roar.BeartypeCallHintReturnViolation)
+
+ParamError = tuple(ParamError)
+ReturnError = tuple(ReturnError)
 
 
 @eqx.filter_jit

test/import_hook_tester_beartype.py

@@ -17,6 +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 equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 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))

test/import_hook_tester_broken_checker.py

@@ -17,6 +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 equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 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))

test/import_hook_tester_transitive/another_file.py

@@ -17,6 +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 equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 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))

test/import_hook_tester_typeguard.py

@@ -17,6 +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 equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 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))

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_import_hook.py

@@ -26,9 +26,8 @@ def test_import_hook_typeguard():
     hook = install_import_hook(
         "test.import_hook_tester_typeguard", ("typeguard", "typechecked")
     )
-    from . import import_hook_tester_typeguard  # noqa: F401
-
-    hook.uninstall()
+    with hook:
+        from . import import_hook_tester_typeguard  # noqa: F401
 
 
 def test_import_hook_beartype():
@@ -40,24 +39,21 @@ def test_import_hook_beartype():
         hook = install_import_hook(
             "test.import_hook_tester_beartype", ("beartype", "beartype")
         )
-        from . import import_hook_tester_beartype  # noqa: F401
-
-        hook.uninstall()
+        with hook:
+            from . import import_hook_tester_beartype  # noqa: F401
 
 
 def test_import_hook_transitive():
     hook = install_import_hook(
         "test.import_hook_tester_transitive", ("typeguard", "typechecked")
     )
-    from . import import_hook_tester_transitive  # noqa: F401
-
-    hook.uninstall()
+    with hook:
+        from . import import_hook_tester_transitive  # noqa: F401
 
 
 def test_import_hook_broken_checker():
     hook = install_import_hook(
         "test.import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist")
     )
-    with pytest.raises(AttributeError):
+    with hook, pytest.raises(AttributeError):
         from . import import_hook_tester_broken_checker  # noqa: F401
-    hook.uninstall()

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

@@ -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()