diff --git a/API.md b/API.md index 104cfa6..7469eec 100644 --- a/API.md +++ b/API.md @@ -62,6 +62,8 @@ Float32[Array, "some_shape"] The array should typically be a `jaxtyping.Array`, which is an alias for `jax.numpy.ndarray`. +`jaxtyping.ArrayLike` is also available, which is an alias for `jax.typing.ArrayLike`. This is a union over JAX arrays and the builtin `bool`/`int`/`float`/`complex`. + But you can use other types as well. `jaxtyping` has support for JAX, NumPy, TensorFlow, and PyTorch, e.g.: ```python Float[np.ndarray, "..."] diff --git a/README.md b/README.md index 6d3dbd9..97de03d 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,10 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]): pip install jaxtyping ``` +Requires Python 3.8+. + +JAX is an optional dependency, required for `jaxtyping.{Array, ArrayLike, PyTree}`. If JAX is not installed then these types will not be available, but you may still use jaxtyping alongside PyTorch/NumPy/etc. + 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 @@ -49,12 +53,6 @@ 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 - -Shape annotations + runtime type checking is inspired by [TorchTyping](https://github.com/patrick-kidger/torchtyping). - -The concise syntax is partially inspired by [etils.array_types](https://github.com/google/etils/tree/main/etils/array_types). - ### Disclaimer This is not an official Google product. diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index b8feb29..33eb5d1 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -18,7 +18,6 @@ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. import typing -import typing_extensions try: @@ -35,6 +34,7 @@ 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 + from jax.typing import ArrayLike as ArrayLike elif has_jax: if getattr(typing, "GENERATING_DOCUMENTATION", False): @@ -42,9 +42,19 @@ elif has_jax: pass Array.__module__ = "builtins" + + class ArrayLike: + pass + + ArrayLike.__module__ = "builtins" else: from jax import Array as Array + try: + from jax.typing import ArrayLike as ArrayLike + except (ModuleNotFoundError, ImportError): + pass + from .array_types import ( AbstractArray as AbstractArray, AbstractDtype as AbstractDtype, diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index f218269..9325def 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -20,8 +20,17 @@ 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 +from typing import ( + Any, + Dict, + List, + Literal, + NoReturn, + Optional, + Tuple, + TYPE_CHECKING, + Union, +) import numpy as np @@ -58,24 +67,72 @@ class _NamedDim: self.name = name self.broadcastable = broadcastable + def __eq__(self, other): + if type(self) is not type(other): + return False + if self.name != other.name: + return False + if self.broadcastable != other.broadcastable: + return False + return True + + def __hash__(self): + return hash((self.name, self.broadcastable)) + class _NamedVariadicDim: def __init__(self, name, broadcastable): self.name = name self.broadcastable = broadcastable + def __eq__(self, other): + if type(self) is not type(other): + return False + if self.name != other.name: + return False + if self.broadcastable != other.broadcastable: + return False + return True + + def __hash__(self): + return hash((self.name, self.broadcastable)) + class _FixedDim: def __init__(self, size, broadcastable): self.size = size self.broadcastable = broadcastable + def __eq__(self, other): + if type(self) is not type(other): + return False + if self.size != other.size: + return False + if self.broadcastable != other.broadcastable: + return False + return True + + def __hash__(self): + return hash((self.size, self.broadcastable)) + class _SymbolicDim: def __init__(self, expr, broadcastable): self.expr = expr self.broadcastable = broadcastable + def __eq__(self, other): + if type(self) is not type(other): + return False + if self.expr != other.expr: + return False + if self.broadcastable != other.broadcastable: + return False + return True + + def __hash__(self): + return hash((self.expr, self.broadcastable)) + _AbstractDimOrVariadicDim = Union[ Literal[_anonymous_dim], @@ -127,6 +184,22 @@ def _check_dims( class _MetaAbstractArray(type): + def __eq__(self, other): + if type(self) is not type(other): + return False + if self.array_type is not other.array_type: + return False + if self.dtypes != other.dtypes: + return False + if self.dims != other.dims: + return False + if self.index_variadic != other.index_variadic: + return False + return True + + def __hash__(self): + return hash((self.array_type, self.dtypes, self.dims, self.index_variadic)) + def __instancecheck__(cls, obj): if not isinstance(obj, cls.array_type): return False @@ -234,6 +307,12 @@ class _MetaAbstractArray(type): assert False +def _check_scalar(dtype, dtypes, dims): + if len(dims) != 0: + return False + return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes) + + class AbstractArray(metaclass=_MetaAbstractArray): array_type: Any dtypes: List[str] @@ -383,31 +462,68 @@ class _MetaAbstractDtype(type): elem = compile(elem, "", "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__}[{type_str}, '{dim_str}']" - elif _array_name_format == "array": - name = type_str + dims = tuple(dims) + + _not_made = object() + + def _make(x): + # Allow Python built-in numeric types. + # TODO: do something more generic than this? Should we _make all types + # that have `shape` and `dtype` attributes or something? + if x is bool: + if _check_scalar("bool", cls.dtypes, dims): + return x + else: + return _not_made + elif x is int: + if _check_scalar("int", cls.dtypes, dims): + return x + else: + return _not_made + elif x is float: + if _check_scalar("float", cls.dtypes, dims): + return x + else: + return _not_made + elif x is complex: + if _check_scalar("complex", cls.dtypes, dims): + return x + else: + return _not_made + try: + type_str = x.__name__ + except AttributeError: + type_str = repr(x) + if _array_name_format == "dtype_and_shape": + name = f"{cls.__name__}[{type_str}, '{dim_str}']" + elif _array_name_format == "array": + name = type_str + else: + raise ValueError( + f"array_name_format {_array_name_format} not recognised" + ) + out = _MetaAbstractArray( + name, + (AbstractArray,), + dict( + array_type=x, + dtypes=cls.dtypes, + dims=dims, + index_variadic=index_variadic, + ), + ) + if getattr(typing, "GENERATING_DOCUMENTATION", False): + out.__module__ = "builtins" + else: + out.__module__ = "jaxtyping" + return out + + if typing.get_origin(array_type) is typing.Union: + out = [_make(x) for x in typing.get_args(array_type)] + out = tuple(x for x in out if x is not _not_made) + out = Union[out] else: - raise ValueError(f"array_name_format {_array_name_format} not recognised") - out = _MetaAbstractArray( - name, - (AbstractArray,), - dict( - array_type=array_type, - dtypes=cls.dtypes, - dims=dims, - index_variadic=index_variadic, - ), - ) - if getattr(typing, "GENERATING_DOCUMENTATION", False): - out.__module__ = "builtins" - else: - out.__module__ = "jaxtyping" + out = _make(array_type) return out @@ -425,7 +541,9 @@ class AbstractDtype(metaclass=_MetaAbstractDtype): dtypes: Union[Literal[_any_dtype], str, List[str]] = cls.dtypes if isinstance(dtypes, str): - dtypes = [dtypes] + dtypes = (dtypes,) + elif dtypes is not _any_dtype: + dtypes = tuple(dtypes) cls.dtypes = dtypes @@ -459,7 +577,8 @@ if TYPE_CHECKING: Annotated as UInt64, ) else: - _bool = "bool_" + _bool = "bool" + _bool_ = "bool_" _uint8 = "uint8" _uint16 = "uint16" _uint32 = "uint32" @@ -502,6 +621,7 @@ else: Complex64 = _make_dtype(_complex64, "Complex64") Complex128 = _make_dtype(_complex128, "Complex128") + bools = [_bool, _bool_] uints = [_uint8, _uint16, _uint32, _uint64] ints = [_int8, _int16, _int32, _int64] floats = [_bfloat16, _float16, _float32, _float64] @@ -510,7 +630,7 @@ else: # We match NumPy's type hierarachy in what types to provide. See the diagram at # https://numpy.org/doc/stable/reference/arrays.scalars.html#scalars - Bool = _make_dtype(_bool, "Bool") + Bool = _make_dtype(bools, "Bool") UInt = _make_dtype(uints, "UInt") Int = _make_dtype(ints, "Int") Integer = _make_dtype(uints + ints, "Integer") diff --git a/setup.py b/setup.py index 27f665d..4f0d2d7 100644 --- a/setup.py +++ b/setup.py @@ -63,7 +63,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Mathematics", ] -python_requires = "~=3.7" +python_requires = "~=3.8" # We use typeguard internally (in a fairly minimal way), but it's not required that # end users make the same choice. diff --git a/test/test_array.py b/test/test_array.py index eb377b3..585fd66 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -17,11 +17,14 @@ # 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 get_args, get_origin, Union + import jax.numpy as jnp import jax.random as jr +import numpy as np import pytest -from jaxtyping import AbstractDtype, Array, Float, Float32, jaxtyped, Shaped +from jaxtyping import AbstractDtype, Array, ArrayLike, Float, Float32, jaxtyped, Shaped from .helpers import ParamError, ReturnError @@ -409,3 +412,54 @@ def test_incomplete_symbolic(typecheck, getkey): x = jr.normal(getkey(), (4,)) with pytest.raises(NameError): foo(x) + + +def test_arraylike(typecheck, getkey): + floatlike1 = Float32[ArrayLike, ""] + floatlike2 = Float[ArrayLike, ""] + floatlike3 = Float32[ArrayLike, "4"] + + assert get_origin(floatlike1) is Union + assert get_origin(floatlike2) is Union + assert get_origin(floatlike3) is Union + assert set(get_args(floatlike1)) == { + Float32[Array, ""], + Float32[np.ndarray, ""], + Float32[np.bool_, ""], + Float32[np.number, ""], + float, + } + assert set(get_args(floatlike2)) == { + Float[Array, ""], + Float[np.ndarray, ""], + Float[np.bool_, ""], + Float[np.number, ""], + float, + } + assert set(get_args(floatlike3)) == { + Float32[Array, "4"], + Float32[np.ndarray, "4"], + Float32[np.bool_, "4"], + Float32[np.number, "4"], + } + + shaped1 = Shaped[ArrayLike, ""] + shaped2 = Shaped[ArrayLike, "4"] + assert get_origin(shaped1) is Union + assert get_origin(shaped2) is Union + assert set(get_args(shaped1)) == { + Shaped[Array, ""], + Shaped[np.ndarray, ""], + Shaped[np.bool_, ""], + Shaped[np.number, ""], + bool, + int, + float, + complex, + } + assert set(get_args(shaped2)) == { + Shaped[Array, "4"], + Shaped[np.ndarray, "4"], + Shaped[np.bool_, "4"], + Shaped[np.number, "4"], + }