Merge pull request #11 from google/v010

Version 0.1.0

.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

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
 

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"

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, "<string>", "eval")
+                elem = _SymbolicDim(elem, broadcastable)
             dims.append(elem)
         if _array_name_format == "dtype_and_shape":
             name = f"{cls.__name__}['{dim_str}']"

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:

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.
 

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"])
 

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)