diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py index e5eaaff..7342701 100644 --- a/jaxtyping/_array_types.py +++ b/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 diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index d23ad16..7a92a9a 100644 --- a/jaxtyping/_decorator.py +++ b/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 != "": diff --git a/test/requirements.txt b/test/requirements.txt index a81b6f9..3670715 100644 --- a/test/requirements.txt +++ b/test/requirements.txt @@ -4,5 +4,6 @@ equinox IPython jaxlib pytest +pytest-asyncio tensorflow typeguard<3 diff --git a/test/test_array.py b/test/test_array.py index 2e99e8c..33c4ab7 100644 --- a/test/test_array.py +++ b/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(): diff --git a/test/test_equals.py b/test/test_equals.py new file mode 100644 index 0000000..b648ff6 --- /dev/null +++ b/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() diff --git a/test/test_generators.py b/test/test_generators.py new file mode 100644 index 0000000..f11fe75 --- /dev/null +++ b/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() diff --git a/test/test_no_jax_dependency.py b/test/test_no_jax_dependency.py index 7b226cd..d1e6f3f 100644 --- a/test/test_no_jax_dependency.py +++ b/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, ) diff --git a/test/test_serialisation.py b/test/test_serialisation.py index f76f1bc..02453b7 100644 --- a/test/test_serialisation.py +++ b/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)