Adding a test for generator support (#171)

* Add a test for generators

* Remove output annotations from decorators

Also guarded torch imports for better compatibility with
requirements.txt

* Add flag to the main meta class to skip the typecheck

* Return to the old solution

* Make async tests work

* Minor adjustments/fixing typos

* Correct Python path for new tests

* Remove some jax-dependent code

* Implement equality for MetaArrays

* Make all Dim variations frozen dataclasses

* Shorten AbstractArray methods

* Final touches

* Removing get_origin use

* Update tests with @jaxtyp

jaxtyping/_array_types.py

@@ -23,6 +23,7 @@ import re
 import sys
 import types
 import typing
+from dataclasses import dataclass
 from typing import Any, Literal, NoReturn, Optional, Union
 
 import numpy as np
@@ -50,7 +51,6 @@ def set_array_name_format(value):
 
 _any_dtype = object()
 
-
 _anonymous_dim = object()
 _anonymous_variadic_dim = object()
 
@@ -61,30 +61,30 @@ class _DimType(enum.Enum):
     symbolic = enum.auto()
 
 
+@dataclass(frozen=True)
 class _NamedDim:
-    def __init__(self, name, broadcastable, treepath):
-        self.name = name
-        self.broadcastable = broadcastable
-        self.treepath = treepath
+    name: str
+    broadcastable: bool
+    treepath: Any
 
 
+@dataclass(frozen=True)
 class _NamedVariadicDim:
-    def __init__(self, name, broadcastable, treepath):
-        self.name = name
-        self.broadcastable = broadcastable
-        self.treepath = treepath
+    name: str
+    broadcastable: bool
+    treepath: Any
 
 
+@dataclass(frozen=True)
 class _FixedDim:
-    def __init__(self, size, broadcastable):
-        self.size = size
-        self.broadcastable = broadcastable
+    size: str
+    broadcastable: bool
 
 
+@dataclass(frozen=True)
 class _SymbolicDim:
-    def __init__(self, elem, broadcastable):
-        self.elem = elem
-        self.broadcastable = broadcastable
+    elem: Any
+    broadcastable: bool
 
 
 _AbstractDimOrVariadicDim = Union[
@@ -147,10 +147,17 @@ def _check_dims(
 
 
 class _MetaAbstractArray(type):
+    _skip_instancecheck: bool = False
+
+    def make_transparent(cls):
+        cls._skip_instancecheck = True
+
     def __instancecheck__(cls, obj: Any) -> bool:
         return cls.__instancecheck_str__(obj) == ""
 
     def __instancecheck_str__(cls, obj: Any) -> str:
+        if cls._skip_instancecheck:
+            return ""
         if not isinstance(obj, cls.array_type):
             return f"this value is not an instance of the underlying array type {cls.array_type}"  # noqa: E501
         if get_treeflatten_memo():
@@ -283,7 +290,24 @@ class _MetaAbstractArray(type):
 @ft.lru_cache(maxsize=None)
 def _make_metaclass(base_metaclass):
     class MetaAbstractArray(_MetaAbstractArray, base_metaclass):
-        pass
+        def _get_props(cls):
+            props_tuple = (
+                cls.index_variadic,
+                cls.dims,
+                cls.array_type,
+                cls.dtypes,
+                cls.dim_str,
+            )
+            return props_tuple
+
+        def __eq__(cls, other):
+            if type(cls) is not type(other):
+                return False
+
+            return cls._get_props() == other._get_props()
+
+        def __hash__(cls):
+            return hash(cls._get_props())
 
     return MetaAbstractArray
 
@@ -314,14 +338,13 @@ class AbstractArray(metaclass=_MetaAbstractArray):
 
 _not_made = object()
 
-
 _union_types = [typing.Union]
 if sys.version_info >= (3, 10):
     _union_types.append(types.UnionType)
 
 
 @ft.lru_cache(maxsize=None)
-def _make_array(array_type, dim_str, dtypes, name):
+def _make_array_cached(array_type, dim_str, dtypes, name):
     if not isinstance(dim_str, str):
         raise ValueError(
             "Shape specification must be a string. Axes should be separated with "
@@ -536,22 +559,33 @@ def _make_array(array_type, dim_str, dtypes, name):
         name = type_str
     else:
         raise ValueError(f"array_name_format {_array_name_format} not recognised")
-    metaclass = _make_metaclass(type(array_type))
-    out = metaclass(
-        name,
-        (array_type, AbstractArray),
-        dict(
-            array_type=array_type,
-            dtypes=dtypes,
-            dims=dims,
-            index_variadic=index_variadic,
-            dim_str=dim_str,
-        ),
-    )
-    if getattr(typing, "GENERATING_DOCUMENTATION", False):
-        out.__module__ = "builtins"
-    else:
-        out.__module__ = "jaxtyping"
+
+    return (array_type, name, dtypes, dims, index_variadic, dim_str)
+
+
+def _make_array(*args, **kwargs):
+    out = _make_array_cached(*args, **kwargs)
+
+    if type(out) is tuple:
+        array_type, name, dtypes, dims, index_variadic, dim_str = out
+        metaclass = _make_metaclass(type(array_type))
+
+        out = metaclass(
+            name,
+            (array_type, AbstractArray),
+            dict(
+                array_type=array_type,
+                dtypes=dtypes,
+                dims=dims,
+                index_variadic=index_variadic,
+                dim_str=dim_str,
+            ),
+        )
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            out.__module__ = "builtins"
+        else:
+            out.__module__ = "jaxtyping"
+
     return out
 
 

jaxtyping/_decorator.py

@@ -26,6 +26,8 @@ import sys
 import warnings
 from typing import Any, get_args, get_origin, get_type_hints, overload
 
+from jaxtyping import AbstractArray
+
 from ._config import config
 from ._errors import AnnotationError, TypeCheckError
 from ._storage import pop_shape_memo, push_shape_memo, shape_str
@@ -309,6 +311,27 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
             # in which case make a best-effort attempt to add shape information for any
             # type errors.
 
+            # we want to detect generators, and ignore return annotations on them,
+            # to avoid issues with O(n) typechecking trying to typecheck yielded values
+            wrp = fn
+            while hasattr(wrp, "__wrapped__"):
+                wrp = wrp.__wrapped__
+
+            if inspect.isgeneratorfunction(wrp) or inspect.isasyncgenfunction(wrp):
+                # recursively parse all the annotations, and mark all the jaxtyping
+                # annotations as not needing instance checks, while still being
+                # visible as original ones for the typechecker
+                def modify_annotation(ann):
+                    if inspect.isclass(ann) and issubclass(ann, AbstractArray):
+                        ann.make_transparent()
+
+                    for sub_ann in get_args(ann):
+                        modify_annotation(sub_ann)
+
+                # just to make sure: check that fn has valid return annotations
+                if hasattr(fn, "__annotations__") and "return" in fn.__annotations__:
+                    modify_annotation(fn.__annotations__["return"])
+
             signature = inspect.signature(fn)
 
             @ft.wraps(fn)
@@ -318,6 +341,7 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
                 try:
                     return fn(*args, **kwargs)
                 except Exception as e:
+                    # add_note api is support from python 3.11+
                     if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
                         shape_info = shape_str(memos)
                         if shape_info != "":

test/requirements.txt

@@ -4,5 +4,6 @@ equinox
 IPython
 jaxlib
 pytest
+pytest-asyncio
 tensorflow
 typeguard<3

test/test_array.py

@@ -25,7 +25,12 @@ import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
 import pytest
-import torch
+
+
+try:
+    import torch
+except ImportError:
+    torch = None
 
 from jaxtyping import (
     AbstractDtype,
@@ -553,7 +558,9 @@ def test_arraylike(typecheck, getkey):
 def test_subclass():
     assert issubclass(Float[Array, ""], Array)
     assert issubclass(Float[np.ndarray, ""], np.ndarray)
-    assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
+
+    if torch is not None:
+        assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
 
 
 def test_ignored_names():

test/test_equals.py

@@ -0,0 +1,35 @@
+from typing import Tuple, Union
+
+import pytest
+
+from jaxtyping import (
+    Array,
+    Float,
+    Float32,
+    Integer,
+    PRNGKeyArray,
+    PyTree,
+    Shaped,
+)
+
+
+@pytest.mark.parametrize(
+    "make_fn",
+    [
+        lambda: Float[Array, "4"],
+        lambda: Float32[Array, ""],
+        lambda: Integer[Array, "1 2 3"],
+        lambda: Shaped[PRNGKeyArray, "2"],
+        lambda: Float[float, "#*shape"],
+        lambda: PyTree[int],
+        lambda: PyTree[Float[Array, ""]],
+        lambda: PyTree[Float32[Array, "*m b c"]],
+        lambda: PyTree[PyTree[Float32[Array, "1 2 b *"]]],
+        lambda: PyTree[Union[str, Float32[Array, "1"]]],
+        lambda: PyTree[
+            Tuple[int, float, Float[Array, ""], PyTree[Union[Float[Array, ""], float]]]
+        ],
+    ],
+)
+def test_equals(make_fn):
+    assert make_fn() == make_fn()

test/test_generators.py

@@ -0,0 +1,88 @@
+from typing import AsyncIterator, Iterator
+
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Array, Float, Shaped
+
+from .helpers import ParamError
+
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+def test_generators_simple(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def gen(x: Float[Array, "*"]) -> Iterator[Float[Array, "*"]]:
+        yield x
+
+    @jaxtyp(typecheck)
+    def foo():
+        next(gen(jnp.zeros(2)))
+        next(gen(jnp.zeros((3, 4))))
+
+    foo()
+
+
+def test_generators_return_no_annotations(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def gen(x: Float[Array, "*"]):
+        yield x
+
+    @jaxtyp(typecheck)
+    def foo():
+        next(gen(jnp.zeros(2)))
+        next(gen(jnp.zeros((3, 4))))
+
+    foo()
+
+
+@pytest.mark.asyncio
+async def test_async_generators_simple(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    async def gen(x: Float[Array, "*"]) -> AsyncIterator[Float[Array, "*"]]:
+        yield x
+
+    @jaxtyp(typecheck)
+    async def foo():
+        async for _ in gen(jnp.zeros(2)):
+            pass
+        async for _ in gen(jnp.zeros((3, 4))):
+            pass
+
+    await foo()
+
+
+def test_generators_dont_modify_same_annotations(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def g(x: Float[Array, "1"]) -> Iterator[Float[Array, "1"]]:
+        yield x
+
+    @jaxtyp(typecheck)
+    def m(x: Float[Array, "1"]) -> Float[Array, "1"]:
+        return x
+
+    with pytest.raises(ParamError):
+        next(g(jnp.zeros(2)))
+    with pytest.raises(ParamError):
+        m(jnp.zeros(2))
+
+
+def test_generators_original_issue(jaxtyp, typecheck):
+    # Effectively the same as https://github.com/patrick-kidger/jaxtyping/issues/91
+    if torch is None:
+        pytest.skip("torch is not available")
+
+    @jaxtyp(typecheck)
+    def g(x: Shaped[torch.Tensor, "*"]) -> Iterator[Shaped[torch.Tensor, "*"]]:
+        yield x
+
+    @jaxtyp(typecheck)
+    def f():
+        next(g(torch.zeros(1)))
+        next(g(torch.zeros(2)))
+
+    f()

test/test_no_jax_dependency.py

@@ -1,9 +1,14 @@
 import subprocess
+import sys
+
+
+_py_path = sys.executable
 
 
 def test_no_jax_dependency():
     result = subprocess.run(
-        "python -c 'import jaxtyping; import sys; sys.exit(\"jax\" in sys.modules)'",
+        f"{_py_path} -c "
+        "'import jaxtyping; import sys; sys.exit(\"jax\" in sys.modules)'",
         shell=True,
     )
     assert result.returncode == 0
@@ -13,7 +18,7 @@ def test_no_jax_dependency():
 # subprocess.)
 def test_meta():
     result = subprocess.run(
-        "python -c 'import jaxtyping; import jax; import sys; "
+        f"{_py_path} -c 'import jaxtyping; import jax; import sys; "
         'sys.exit("jax" in sys.modules)\'',
         shell=True,
     )

test/test_serialisation.py

@@ -1,16 +1,25 @@
 import cloudpickle
 import numpy as np
-import torch
+
+
+try:
+    import torch
+except ImportError:
+    torch = None
 
 from jaxtyping import AbstractArray, Array, Shaped
 
 
 def test_pickle():
     x = cloudpickle.dumps(Shaped[Array, ""])
-    y = cloudpickle.dumps(AbstractArray)
-    z = cloudpickle.dumps(Shaped[np.ndarray, ""])
-    w = cloudpickle.dumps(Shaped[torch.Tensor, ""])
     cloudpickle.loads(x)
+
+    y = cloudpickle.dumps(AbstractArray)
     cloudpickle.loads(y)
+
+    z = cloudpickle.dumps(Shaped[np.ndarray, ""])
     cloudpickle.loads(z)
-    cloudpickle.loads(w)
+
+    if torch is not None:
+        w = cloudpickle.dumps(Shaped[torch.Tensor, ""])
+        cloudpickle.loads(w)