diff --git a/.flake8 b/.flake8 index 7b18abd..8026b92 100644 --- a/.flake8 +++ b/.flake8 @@ -1,4 +1,4 @@ [flake8] -max-line-length = 120 +max-line-length = 88 ignore = W291,W293,W503,W504,E123,E126,E203,E402,E701,E731,F722 per-file-ignores = __init__.py: F401 diff --git a/API.md b/API.md index bcf8ce1..d82e65d 100644 --- a/API.md +++ b/API.md @@ -6,22 +6,29 @@ Each array is denoted by a type `dtype[shape]`, such as `f32["batch channels"]`. ### Shape -The shape should be a string of space-separated symbols, such as "a b c d". Each symbol can be: +The shape should be a string of space-separated symbols, such as "a b c d". Each symbol can be either an: - `int`: fixed-size axis, e.g. `f32["28 28"]`. - `str`: variable-size axis, e.g. `f32["channels"]`. -- `_`: anonymous axis, e.g. `f32["batch channels _ _"]`. -- `...`: anonymous zero or more axes, e.g. `f32["... c h w"]` -- `*name`: zero or more variable-size axes, e.g. `f32["*batch c h w"]` -- Append `#` to a dimension size to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable. +- A symbolic expression (without spaces!) in terms of other variable-size axes, e.g. `def remove_last(x: f32["dim"]) -> f32["dim-1"]`. -When calling a function, variable-size axes will be matched up across all arguments and checked for consistency. (See [runtime type checking](#runtime-type-checking) below.) +When calling a function, variable-size axes and symbolic axes will be matched up across all arguments and checked for consistency. (See [runtime type checking](#runtime-type-checking) below.) + +In addition some modifiers can be applied: +- Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `f32["*batch c h w"]` will match zero or more batch axes. +- Prepend `#` to a dimension to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g. `add(x: f32["#foo"], y: f32["#foo"]) -> f32["#foo"]`. +- Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `f32["b c _ _"]`. + +The order of these modifiers does not matter. + +As a special case: +- `...`: anonymous zero or more axes (equivalent to `*_`) e.g. `f32["... c h w"]` Some notes: - To denote a scalar shape use `""`, e.g. `f32[""]`. - To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `f32["..."]`. -- You cannot have multiple variadic axes, i.e. you can only use `...` or `*name` at most once in each array. -- An example of broadcasting in one dimension: `add(x: f32["foo#"], y: f32["foo#"]) -> f32["foo#"]`. -- An example of broadcasting multiple dimensions: `add(x: f32["*foo#"], y: f32["*foo#"]) -> f32["*foo#"]`. +- You cannot have more than one use of multiple-axes, i.e. you can only use `...` or `*name` at most once in each array. +- An example of broadcasting multiple dimensions: `add(x: f32["*#foo"], y: f32["*#foo"]) -> f32["*#foo"]`. +- A symbolic expression cannot be evaluated unless all of the axes sizes it refers to have already been processed. In practice this usually means that they should only be used in annotations for the return type, and only use axes declared in the arguments. ### Dtype diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index bf4a5ad..41a8428 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -51,4 +51,4 @@ from .import_hook import install_import_hook from .pytree_type import PyTree -__version__ = "0.0.2" +__version__ = "0.1.0" diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index 348ee06..8b62f22 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -17,11 +17,13 @@ # 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 enum import functools as ft from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union from typing_extensions import Literal import jax.numpy as jnp +import numpy as np from .decorator import storage @@ -40,10 +42,17 @@ def set_array_name_format(value): _any_dtype = object() + _anonymous_dim = object() _anonymous_variadic_dim = object() +class _DimType(enum.Enum): + named = enum.auto() + fixed = enum.auto() + symbolic = enum.auto() + + class _NamedDim: def __init__(self, name, broadcastable): self.name = name @@ -62,21 +71,28 @@ class _FixedDim: self.broadcastable = broadcastable +class _SymbolicDim: + def __init__(self, expr, broadcastable): + self.expr = expr + self.broadcastable = broadcastable + + _AbstractDimOrVariadicDim = Union[ Literal[_anonymous_dim], Literal[_anonymous_variadic_dim], _NamedDim, _NamedVariadicDim, _FixedDim, + _SymbolicDim, ] -_AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim] +_AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim, _SymbolicDim] def _check_dims( cls_dims: List[_AbstractDim], obj_shape: Tuple[int], - memo: Dict[str, Union[int, Tuple[int]]], -): + single_memo: Dict[str, int], +) -> bool: assert len(cls_dims) == len(obj_shape) for cls_dim, obj_size in zip(cls_dims, obj_shape): if cls_dim is _anonymous_dim: @@ -86,12 +102,24 @@ def _check_dims( elif type(cls_dim) is _FixedDim: if cls_dim.size != obj_size: return False + elif type(cls_dim) is _SymbolicDim: + try: + eval_size = eval(cls_dim.expr, single_memo) + except NameError as e: + raise NameError( + f"Cannot process symbolic dimension '{cls_dim.expr}' as some " + "dimension names have not been processed. In practice you should " + "usually only use symbolic dimensions in annotations for return " + "types, referring only to dimensions annotated for arguments." + ) from e + if eval_size != obj_size: + return False else: assert type(cls_dim) is _NamedDim try: - cls_size = memo[cls_dim.name] + cls_size = single_memo[cls_dim.name] except KeyError: - memo[cls_dim.name] = obj_size + single_memo[cls_dim.name] = obj_size else: if cls_size != obj_size: return False @@ -110,26 +138,41 @@ class _MetaAbstractArray(type): # `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. - memo = {} + 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. - memo = storage.memo_stack[-1].copy() + 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, memo): + 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: - storage.memo_stack[-1] = memo + storage.memo_stack[-1] = ( + single_memo, + variadic_memo, + variadic_broadcast_memo, + ) return True else: return False - def _check_shape(cls, obj, memo): + def _check_shape( + cls, + obj, + single_memo: Dict[str, int], + variadic_memo: Dict[str, Tuple[int, ...]], + variadic_broadcast_memo: Dict[str, List[Tuple[int, ...]]], + ): if cls.index_variadic is None: if obj.ndim != len(cls.dims): return False - return _check_dims(cls.dims, obj.shape, memo) + return _check_dims(cls.dims, obj.shape, single_memo) else: if obj.ndim < len(cls.dims) - 1: return False @@ -137,34 +180,42 @@ class _MetaAbstractArray(type): j = -(len(cls.dims) - i - 1) if j == 0: j = None - if not _check_dims(cls.dims[:i], obj.shape[:i], memo): + if not _check_dims(cls.dims[:i], obj.shape[:i], single_memo): return False - if j is not None and not _check_dims(cls.dims[j:], obj.shape[j:], memo): + if j is not None and not _check_dims( + cls.dims[j:], obj.shape[j:], single_memo + ): return False variadic_dim = cls.dims[i] - if variadic_dim is not _anonymous_variadic_dim: + if variadic_dim is _anonymous_variadic_dim: + return True + else: + assert type(variadic_dim) is _NamedVariadicDim variadic_name = variadic_dim.name try: - variadic_shape = memo[variadic_name] + if variadic_dim.broadcastable: + variadic_shapes = variadic_broadcast_memo[variadic_name] + else: + variadic_shape = variadic_memo[variadic_name] except KeyError: - memo[variadic_name] = obj.shape[i:j] + if variadic_dim.broadcastable: + variadic_broadcast_memo[variadic_name] = [obj.shape[i:j]] + else: + variadic_memo[variadic_name] = obj.shape[i:j] + return True else: if variadic_dim.broadcastable: - new_variadic_shape = [] - obj_shape = obj.shape[i:j] - if len(variadic_shape) != len(obj_shape): - return False - for old_size, new_size in zip(variadic_shape, obj_shape): - if old_size == 1: - new_variadic_shape.append(new_size) - else: - if new_size != 1 and old_size != new_size: - return False - new_variadic_shape.append(old_size) - memo[variadic_name] = tuple(new_variadic_shape) + new_shape = obj.shape[i:j] + for existing_shape in variadic_shapes: + try: + np.broadcast_shapes(new_shape, existing_shape) + except ValueError: + return False + variadic_shapes.append(new_shape) + return True else: return variadic_shape == obj.shape[i:j] - return True + assert False class AbstractArray(metaclass=_MetaAbstractArray): @@ -185,7 +236,8 @@ class _MetaAbstractDtype(type): def __getitem__(cls, dim_str: str) -> _MetaAbstractArray: if not isinstance(dim_str, str): raise ValueError( - "Shape specification must be a string. Axes should be separated with spaces." + "Shape specification must be a string. Axes should be separated with " + "spaces." ) dims = [] index_variadic = None @@ -195,29 +247,117 @@ class _MetaAbstractDtype(type): raise ValueError( "Dimensions should be separated with spaces, not commas" ) - broadcastable = False if elem.endswith("#"): - broadcastable = True - elem = elem[:-1] - try: - elem = int(elem) - except ValueError: - if elem == "_": - elem = _anonymous_dim - elif elem == "...": - if index_variadic is not None: - raise ValueError("Cannot have multiple variadic dimensions") - index_variadic = index - elem = _anonymous_variadic_dim - elif elem[0] == "*": - if index_variadic is not None: - raise ValueError("Cannot have multiple variadic dimensions") - index_variadic = index - elem = _NamedVariadicDim(elem[1:], broadcastable) - else: - elem = _NamedDim(elem, broadcastable) + raise ValueError( + "As of jaxtyping v0.1.0, broadcastable dimensions are now denoted " + "with a # at the start, rather than at the end" + ) + + if "..." in elem: + if elem != "...": + raise ValueError( + "Anonymous multiple dimension '...' must be used on its own; " + f"got {elem}" + ) + broadcastable = False + variadic = True + anonymous = True + dim_type = _DimType.named else: + broadcastable = False + variadic = False + anonymous = False + while True: + if len(elem) == 0: + # This branch needed as just `_` is valid + break + first_char = elem[0] + if first_char == "#": + if broadcastable: + raise ValueError( + "Do not use # twice to denote broadcastability, e.g. " + "`##foo` is not allowed" + ) + broadcastable = True + elem = elem[1:] + elif first_char == "*": + if variadic: + raise ValueError( + "Do not use * twice to denote accepting multiple " + "dimensions, e.g. `**foo` is not allowed" + ) + variadic = True + elem = elem[1:] + elif first_char == "_": + if anonymous: + raise ValueError( + "Do not use _ twice to denote anonymity, e.g. `__foo` " + "is not allowed" + ) + anonymous = True + elem = elem[1:] + else: + break + try: + elem = int(elem) + except ValueError: + if len(elem) == 0 or elem.isidentifier(): + dim_type = _DimType.named + else: + dim_type = _DimType.symbolic + else: + dim_type = _DimType.fixed + + if variadic: + if index_variadic is not None: + raise ValueError( + "Cannot use multiple-dimension specifiers (`*name` or `...`) " + "more than once" + ) + index_variadic = index + + if dim_type is _DimType.fixed: + if variadic: + raise ValueError( + "Cannot have a fixed axis bind to multiple dimensions, e.g. " + "`*4` is not allowed" + ) + if anonymous: + raise ValueError( + "Cannot have a fixed axis be anonymous, e.g. `_4` is not " + "allowed" + ) elem = _FixedDim(elem, broadcastable) + elif dim_type is _DimType.named: + if anonymous: + if broadcastable: + raise ValueError( + "Cannot have a dimension be both anonymous and " + "broadcastable, e.g. `#_` is not allowed" + ) + if variadic: + elem = _anonymous_variadic_dim + else: + elem = _anonymous_dim + else: + if variadic: + elem = _NamedVariadicDim(elem, broadcastable) + else: + elem = _NamedDim(elem, broadcastable) + else: + assert dim_type is _DimType.symbolic + if anonymous: + raise ValueError( + "Cannot have a symbolic dimension be anonymous, e.g. " + "`_foo+bar` is not allowed" + ) + if variadic: + raise ValueError( + "Cannot have symbolic multiple-dimensions, e.g. " + "`*foo+bar` is not allowed" + ) + elem = compile(elem, "", "eval") + elem = _SymbolicDim(elem, broadcastable) dims.append(elem) if _array_name_format == "dtype_and_shape": name = f"{cls.__name__}['{dim_str}']" diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py index 5d34016..3099007 100644 --- a/jaxtyping/decorator.py +++ b/jaxtyping/decorator.py @@ -28,8 +28,7 @@ storage.memo_stack = [] def jaxtyped(fn): @ft.wraps(fn) def wrapper(*args, **kwargs): - memo = {} - storage.memo_stack.append(memo) + storage.memo_stack.append(({}, {}, {})) try: return fn(*args, **kwargs) finally: diff --git a/jaxtyping/import_hook.py b/jaxtyping/import_hook.py index 31f74e8..a8dce1a 100644 --- a/jaxtyping/import_hook.py +++ b/jaxtyping/import_hook.py @@ -31,19 +31,20 @@ # # 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: +# 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 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. +# 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. # # --------- @@ -148,7 +149,8 @@ class _JaxtypingLoader(SourceFileLoader): ) def exec_module(self, module): - # Use a custom optimization marker – the import lock should make this monkey patch safe + # Use a custom optimization marker - the import lock should make this monkey + # patch safe with patch( "importlib._bootstrap_external.cache_from_source", _optimized_cache_from_source, @@ -216,7 +218,8 @@ class ImportHookManager: def install_import_hook( modules: Iterable[str], typechecker: Optional[Tuple[str, str]] ) -> ImportHookManager: - """Automatically apply `@jaxtyped`, and optionally a type checker, to all classes and functions. + """Automatically apply `@jaxtyped`, and optionally a type checker, to all classes + and functions. It will only be applied to modules loaded **after** this hook has been installed. diff --git a/setup.py b/setup.py index 1a97619..37b29b3 100644 --- a/setup.py +++ b/setup.py @@ -28,7 +28,8 @@ _here = pathlib.Path(__file__).resolve().parent name = "jaxtyping" -# for simplicity we actually store the version in the __version__ attribute in the source +# for simplicity we actually store the version in the __version__ attribute in the +# source with open(_here / name / "__init__.py") as f: meta_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", f.read(), re.M) if meta_match: @@ -40,7 +41,10 @@ author = "Patrick Kidger" author_email = "contact@kidger.site" -description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees." +description = ( + "Type annotations and runtime checking for shape and dtype of JAX " + "arrays, and PyTrees." +) with open(_here / "README.md", "r") as f: readme = f.read() @@ -63,7 +67,12 @@ 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. -install_requires = ["jax>=0.3.4", "typeguard>=2.13.3", "typing_extensions>=4.2.0"] +install_requires = [ + "jax>=0.3.4", + "numpy>=1.20.0", + "typeguard>=2.13.3", + "typing_extensions>=4.2.0", +] entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"]) diff --git a/test/test_array.py b/test/test_array.py index e6de38b..7f94f24 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -227,7 +227,7 @@ def test_anonymous_variadic(typecheck, getkey): def test_broadcast_fixed(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32["4#"]): + def g(x: f32["#4"]): pass g(jr.normal(getkey(), (4,))) @@ -240,7 +240,7 @@ def test_broadcast_fixed(typecheck, getkey): def test_broadcast_named(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[" foo#"], y: f32[" foo#"]): + def g(x: f32[" #foo"], y: f32[" #foo"]): pass a = jr.normal(getkey(), (3,)) @@ -264,7 +264,7 @@ def test_broadcast_named(typecheck, getkey): def test_broadcast_variadic_named(typecheck, getkey): @jaxtyped @typecheck - def g(x: f32[" *foo#"], y: f32[" *foo#"]): + def g(x: f32[" *#foo"], y: f32[" *#foo"]): pass a = jr.normal(getkey(), (3,)) @@ -283,12 +283,11 @@ def test_broadcast_variadic_named(typecheck, getkey): g(b, b) g(c, c) g(d, d) + g(b, c) with pytest.raises(ParamError): g(a, b) with pytest.raises(ParamError): g(a, c) - with pytest.raises(ParamError): - g(b, c) with pytest.raises(ParamError): g(a, b) with pytest.raises(ParamError): @@ -296,26 +295,20 @@ def test_broadcast_variadic_named(typecheck, getkey): g(a, j) g(b, j) - with pytest.raises(ParamError): - g(c, j) - with pytest.raises(ParamError): - g(d, j) - with pytest.raises(ParamError): - g(b, k) + g(c, j) + g(d, j) + g(b, k) g(c, k) with pytest.raises(ParamError): g(d, k) with pytest.raises(ParamError): g(c, l) g(d, l) - with pytest.raises(ParamError): - g(a, m) + g(a, m) g(c, m) g(d, m) - with pytest.raises(ParamError): - g(a, n) - with pytest.raises(ParamError): - g(b, n) + g(a, n) + g(b, n) with pytest.raises(ParamError): g(c, n) with pytest.raises(ParamError): @@ -327,6 +320,54 @@ def test_broadcast_variadic_named(typecheck, getkey): g(o, a) -def test_no_commas(typecheck, getkey): +def test_no_commas(): with pytest.raises(ValueError): f32["foo, bar"] + + +def test_symbolic(typecheck, getkey): + @jaxtyped + @typecheck + def make_slice(x: f32[" dim"]) -> f32[" dim-1"]: + return x[1:] + + @jaxtyped + @typecheck + def cat(x: f32[" dim"]) -> f32[" 2*dim"]: + return jnp.concatenate([x, x]) + + @jaxtyped + @typecheck + def bad_make_slice(x: f32[" dim"]) -> f32[" dim-1"]: + return x + + @jaxtyped + @typecheck + def bad_cat(x: f32[" dim"]) -> f32[" 2*dim"]: + return jnp.concatenate([x, x, x]) + + x = jr.normal(getkey(), (5,)) + assert make_slice(x).shape == (4,) + assert cat(x).shape == (10,) + + y = jr.normal(getkey(), (3, 4)) + with pytest.raises(ParamError): + make_slice(y) + with pytest.raises(ParamError): + cat(y) + + with pytest.raises(ReturnError): + bad_make_slice(x) + with pytest.raises(ReturnError): + bad_cat(x) + + +def test_incomplete_symbolic(typecheck, getkey): + @jaxtyped + @typecheck + def foo(x: f32[" 2*dim"]): + pass + + x = jr.normal(getkey(), (4,)) + with pytest.raises(NameError): + foo(x)