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

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)