Added support for treepath-dependent sizes.

docs/api/array.md

@@ -24,6 +24,7 @@ In addition some modifiers can be applied:
     `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#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. `"b c _ _"`.
 - Documentation-only names (i.e. they're ignored by jaxtyping) can be handled by prepending a name followed by `=` e.g. `Float[Array, "rows=4 cols=3"]`.
+- Prepend `?` to a dimension to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree.md).)
 
 When using multiple modifiers, their order does not matter.
 

docs/api/pytree.md

@@ -11,4 +11,36 @@
 
 ---
 
+## Path-dependent shapes
+
+The prefix `?` may be used to indicate that the axis size can depend on which leaf of a PyTree the array is at. For example:
+```python
+def f(
+    x: PyTree[Shaped[Array, "?foo"], "T"],
+    y: PyTree[Shaped[Array, "?foo"], "T"],
+):
+    pass
+```
+The above demands that `x` and `y` have matching PyTree structures (due to the `T` annotation), and that their leaves must all be one-dimensional arrays, *and that the corresponding pairs of leaves in `x` and `y` must have the same size as each other*.
+
+Thus the following is allowed:
+```python
+x0 = jnp.arange(3)
+x1 = jnp.arange(5)
+
+y0 = jnp.arange(3) + 1
+y1 = jnp.arange(5) + 1
+
+f((x0, x1), (y0, y1))  # x0 matches y0, and x1 matches y1. All good!
+```
+
+But this is not:
+```python
+f((x1, x1), (y0, y1))  # x1 does not have a size matching y0!
+```
+
+Internally, all that is happening is that `foo` is replaced with `0foo` for the first leaf, `1foo` for the next leaf, etc., so that each leaf gets a unique version of the name.
+
+---
+
 Note that `jaxtyping.{PyTree, PyTreeDef}` are only available if JAX has been installed.

jaxtyping/_array_types.py

@@ -27,7 +27,7 @@ from typing import Any, Literal, NoReturn, Optional, Union
 
 import numpy as np
 
-from ._decorator import storage_get, storage_set
+from ._storage import get_shape_memo, get_treepath_memo, set_shape_memo
 
 
 try:
@@ -68,15 +68,17 @@ class _DimType(enum.Enum):
 
 
 class _NamedDim:
-    def __init__(self, name, broadcastable):
+    def __init__(self, name, broadcastable, treepath):
         self.name = name
         self.broadcastable = broadcastable
+        self.treepath = treepath
 
 
 class _NamedVariadicDim:
-    def __init__(self, name, broadcastable):
+    def __init__(self, name, broadcastable, treepath):
         self.name = name
         self.broadcastable = broadcastable
+        self.treepath = treepath
 
 
 class _FixedDim:
@@ -133,10 +135,14 @@ def _check_dims(
                 return False
         else:
             assert type(cls_dim) is _NamedDim
+            if cls_dim.treepath:
+                name = get_treepath_memo() + cls_dim.name
+            else:
+                name = cls_dim.name
             try:
-                cls_size = single_memo[cls_dim.name]
+                cls_size = single_memo[name]
             except KeyError:
-                single_memo[cls_dim.name] = obj_size
+                single_memo[name] = obj_size
             else:
                 if cls_size != obj_size:
                     return False
@@ -198,12 +204,19 @@ class _MetaAbstractArray(type):
             if not in_dtypes:
                 return False
 
-        single_memo, variadic_memo, pytree_memo = storage_get()
-        if cls._check_shape(obj, single_memo, variadic_memo):
-            # We update the memo every time we successfully pass a shape check
-            storage_set(single_memo, variadic_memo, pytree_memo)
+        single_memo, variadic_memo, pytree_memo = get_shape_memo()
+        single_memo_bak = single_memo.copy()
+        variadic_memo_bak = variadic_memo.copy()
+        pytree_memo_bak = pytree_memo.copy()
+        try:
+            check = cls._check_shape(obj, single_memo, variadic_memo)
+        except Exception:
+            set_shape_memo(single_memo_bak, variadic_memo_bak, pytree_memo_bak)
+            raise
+        if check:
             return True
         else:
+            set_shape_memo(single_memo_bak, variadic_memo_bak, pytree_memo_bak)
             return False
 
     def _check_shape(
@@ -234,7 +247,10 @@ class _MetaAbstractArray(type):
                 return True
             else:
                 assert type(variadic_dim) is _NamedVariadicDim
-                name = variadic_dim.name
+                if variadic_dim.treepath:
+                    name = get_treepath_memo() + variadic_dim.name
+                else:
+                    name = variadic_dim.name
                 broadcastable = variadic_dim.broadcastable
                 try:
                     prev_broadcastable, prev_shape = variadic_memo[name]
@@ -338,11 +354,13 @@ def _make_array(array_type, dim_str, dtypes, name):
             broadcastable = False
             variadic = True
             anonymous = True
+            treepath = False
             dim_type = _DimType.named
         else:
             broadcastable = False
             variadic = False
             anonymous = False
+            treepath = False
             while True:
                 if len(elem) == 0:
                     # This branch needed as just `_` is valid
@@ -372,6 +390,14 @@ def _make_array(array_type, dim_str, dtypes, name):
                         )
                     anonymous = True
                     elem = elem[1:]
+                elif first_char == "?":
+                    if treepath:
+                        raise ValueError(
+                            "Do not use ? twice to denote dependence on location "
+                            "within a PyTree, e.g. `??foo` is not allowed"
+                        )
+                    treepath = True
+                    elem = elem[1:]
                 # Allow e.g. `foo=4` as an alternate syntax for just `4`, so that one
                 # can write e.g. `Float[Array, "rows=3 cols=4"]`
                 elif elem.count("=") == 1:
@@ -392,7 +418,7 @@ def _make_array(array_type, dim_str, dtypes, name):
             if index_variadic is not None:
                 raise ValueError(
                     "Cannot use multiple-dimension specifiers (`*name` or `...`) "
-                    "more than once"
+                    "more than once."
                 )
             index_variadic = index
 
@@ -400,11 +426,16 @@ def _make_array(array_type, dim_str, dtypes, name):
             if variadic:
                 raise ValueError(
                     "Cannot have a fixed axis bind to multiple dimensions, e.g. "
-                    "`*4` is not allowed"
+                    "`*4` is not allowed."
                 )
             if anonymous:
                 raise ValueError(
-                    "Cannot have a fixed axis be anonymous, e.g. `_4` is not " "allowed"
+                    "Cannot have a fixed axis be anonymous, e.g. `_4` is not allowed."
+                )
+            if treepath:
+                raise ValueError(
+                    "Cannot have a fixed axis have tree-path dependence, e.g. `?4` is "
+                    "not allowed."
                 )
             elem = _FixedDim(elem, broadcastable)
         elif dim_type is _DimType.named:
@@ -412,7 +443,7 @@ def _make_array(array_type, dim_str, dtypes, name):
                 if broadcastable:
                     raise ValueError(
                         "Cannot have a dimension be both anonymous and "
-                        "broadcastable, e.g. `#_` is not allowed"
+                        "broadcastable, e.g. `#_` is not allowed."
                     )
                 if variadic:
                     elem = _anonymous_variadic_dim
@@ -420,9 +451,9 @@ def _make_array(array_type, dim_str, dtypes, name):
                     elem = _anonymous_dim
             else:
                 if variadic:
-                    elem = _NamedVariadicDim(elem, broadcastable)
+                    elem = _NamedVariadicDim(elem, broadcastable, treepath)
                 else:
-                    elem = _NamedDim(elem, broadcastable)
+                    elem = _NamedDim(elem, broadcastable, treepath)
         else:
             assert dim_type is _DimType.symbolic
             if anonymous:
@@ -435,6 +466,11 @@ def _make_array(array_type, dim_str, dtypes, name):
                     "Cannot have symbolic multiple-dimensions, e.g. "
                     "`*foo+bar` is not allowed"
                 )
+            if treepath:
+                raise ValueError(
+                    "Cannot have a symbolic dimensions with tree-path dependence, e.g. "
+                    "`?foo+bar` is not allowed"
+                )
             elem_string = elem
             elem = compile(elem, "<string>", "eval")
             elem = _SymbolicDim(elem, broadcastable, elem_string)

jaxtyping/_decorator.py

@@ -21,7 +21,6 @@ import dataclasses
 import functools as ft
 import inspect
 import textwrap
-import threading
 import types
 import weakref
 from typing import get_args, get_origin
@@ -35,33 +34,7 @@ else:
     traceback_util.register_exclusion(__file__)
 
 
-_storage = threading.local()
-
-
-def _no_temp_memo():
-    return hasattr(_storage, "memo_stack") and len(_storage.memo_stack) != 0
-
-
-def storage_get():
-    if _no_temp_memo():
-        single_memo, variadic_memo, pytree_memo = _storage.memo_stack[-1]
-        # Make a copy so we don't mutate the original memo during the shape check.
-        single_memo = single_memo.copy()
-        variadic_memo = variadic_memo.copy()
-        pytree_memo = pytree_memo.copy()
-    else:
-        # `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.
-        single_memo = {}
-        variadic_memo = {}
-        pytree_memo = {}
-    return single_memo, variadic_memo, pytree_memo
-
-
-def storage_set(single_memo, variadic_memo, pytree_memo):
-    if _no_temp_memo():
-        _storage.memo_stack[-1] = single_memo, variadic_memo, pytree_memo
+from ._storage import pop_shape_memo, push_shape_memo
 
 
 _jaxtyped_fns = weakref.WeakSet()
@@ -150,15 +123,11 @@ def jaxtyped(fn):
 
         @ft.wraps(fn)
         def wrapped_fn(*args, **kwargs):
-            try:
-                memo_stack = _storage.memo_stack
-            except AttributeError:
-                memo_stack = _storage.memo_stack = []
-            memo_stack.append(({}, {}, {}))
+            push_shape_memo()
             try:
                 return fn(*args, **kwargs)
             finally:
-                memo_stack.pop()
+                pop_shape_memo()
 
         _jaxtyped_fns.add(wrapped_fn)
         return wrapped_fn
@@ -166,14 +135,10 @@ def jaxtyped(fn):
 
 class _JaxtypingContext:
     def __enter__(self):
-        try:
-            memo_stack = _storage.memo_stack
-        except AttributeError:
-            memo_stack = _storage.memo_stack = []
-        memo_stack.append(({}, {}, {}))
+        push_shape_memo()
 
     def __exit__(self, exc_type, exc_value, exc_tb):
-        _storage.memo_stack.pop()
+        pop_shape_memo()
 
 
 @jaxtyped

jaxtyping/_pytree_type.py

@@ -24,7 +24,12 @@ from typing import Generic, TypeVar
 import jax.tree_util as jtu
 import typeguard
 
-from ._decorator import storage_get, storage_set
+from ._storage import (
+    clear_treepath_memo,
+    get_shape_memo,
+    set_shape_memo,
+    set_treepath_memo,
+)
 
 
 _T = TypeVar("_T")
@@ -47,6 +52,22 @@ class _MetaPyTree(type):
         if not hasattr(cls, "leaftype"):
             return True  # Just `isinstance(x, PyTree)`
 
+        single_memo, variadic_memo, pytree_memo = get_shape_memo()
+        single_memo_bak = single_memo.copy()
+        variadic_memo_bak = variadic_memo.copy()
+        pytree_memo_bak = pytree_memo.copy()
+        try:
+            out = cls._check(obj, pytree_memo)
+        except Exception:
+            set_shape_memo(single_memo_bak, variadic_memo_bak, pytree_memo_bak)
+            raise
+        if out:
+            return True
+        else:
+            set_shape_memo(single_memo_bak, variadic_memo_bak, pytree_memo_bak)
+            return False
+
+    def _check(cls, obj, pytree_memo):
         # We could use `isinstance` here but that would fail for more complicated
         # types, e.g. PyTree[tuple[int]]. So at least internally we make a particular
         # choice of typechecker.
@@ -57,16 +78,20 @@ class _MetaPyTree(type):
         def accepts_leaftype(x: cls.leaftype):
             pass
 
-        def is_leaftype(x):
+        def is_leaftype(x, new_scope=True):
+            if new_scope and cls.structure is not None:
+                set_treepath_memo(-1, "")
             try:
                 accepts_leaftype(x)
             except _TypeCheckError:
                 return False
             else:
                 return True
+            finally:
+                if new_scope and cls.structure is not None:
+                    clear_treepath_memo()
 
         leaves, structure = jtu.tree_flatten(obj, is_leaf=is_leaftype)
-        single_memo, variadic_memo, pytree_memo = storage_get()
         if cls.structure is not None:
             if cls.structure.isidentifier():
                 try:
@@ -122,10 +147,16 @@ class _MetaPyTree(type):
                 else:
                     if structure != named_structure:
                         return False
-        for leaf_index, leaf in enumerate(leaves):
-            if not is_leaftype(leaf):
-                return False
-        storage_set(single_memo, variadic_memo, pytree_memo)
+
+        try:
+            for leaf_index, leaf in enumerate(leaves):
+                if cls.structure is not None:
+                    set_treepath_memo(leaf_index, cls.structure)
+                if not is_leaftype(leaf, new_scope=False):
+                    return False
+                clear_treepath_memo()
+        finally:
+            clear_treepath_memo()
         return True
 
     # Can't return a generic (e.g. _FakePyTree[item]) because generic aliases don't do

jaxtyping/_storage.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2022 Google LLC
+#
+# 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:
+#
+# 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.
+
+import threading
+
+
+_shape_storage = threading.local()
+
+
+def _has_shape_memo():
+    return hasattr(_shape_storage, "memo_stack") and len(_shape_storage.memo_stack) != 0
+
+
+def get_shape_memo():
+    if _has_shape_memo():
+        single_memo, variadic_memo, pytree_memo = _shape_storage.memo_stack[-1]
+    else:
+        # `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.
+        single_memo = {}
+        variadic_memo = {}
+        pytree_memo = {}
+    return single_memo, variadic_memo, pytree_memo
+
+
+def set_shape_memo(single_memo, variadic_memo, pytree_memo) -> None:
+    if _has_shape_memo():
+        _shape_storage.memo_stack[-1] = single_memo, variadic_memo, pytree_memo
+
+
+def push_shape_memo() -> None:
+    try:
+        memo_stack = _shape_storage.memo_stack
+    except AttributeError:
+        # Can't be done when `_stack_storage` is created for reasons I forget.
+        memo_stack = _shape_storage.memo_stack = []
+    memo_stack.append(({}, {}, {}))
+
+
+def pop_shape_memo() -> None:
+    _shape_storage.memo_stack.pop()
+
+
+_treepath_storage = threading.local()
+
+
+def clear_treepath_memo() -> None:
+    _treepath_storage.value = None
+
+
+def set_treepath_memo(index: int, structure: str) -> None:
+    if hasattr(_treepath_storage, "value") and _treepath_storage.value is not None:
+        raise ValueError(
+            "Cannot typecheck annotations of the form "
+            "`PyTree[PyTree[Shaped[Array, '?foo'], 'T'], 'S']` as it is ambiguous "
+            "which PyTree the `?` annotation refers to."
+        )
+    _treepath_storage.value = str(index) + structure
+
+
+def get_treepath_memo() -> str:
+    if not hasattr(_treepath_storage, "value") or _treepath_storage.value is None:
+        raise ValueError(
+            "Cannot use `?` annotations, e.g. `Shaped[Array, '?foo']`, except "
+            "when contained with structured `PyTree` annotations, e.g. "
+            "`PyTree[Shaped[Array, '?foo'], 'T']`."
+        )
+    return _treepath_storage.value

test/test_pytree.py

@@ -25,7 +25,8 @@ import jax.numpy as jnp
 import jax.random as jr
 import pytest
 
-from jaxtyping import Float, jaxtyped, PyTree
+import jaxtyping
+from jaxtyping import Array, Float, jaxtyped, PyTree
 
 from .helpers import make_mlp, ParamError
 
@@ -270,3 +271,84 @@ def test_structure_compose(typecheck):
         g((1, 2), {"a": 3}, {"a": ("hi", "bye")})
 
     g((1, 2), {"a": 3}, ({"a": "hi"}, {"a": "bye"}))
+
+
+@pytest.mark.parametrize("variadic", (False, True))
+def test_treepath_dependence_function(variadic, typecheck, getkey):
+    if variadic:
+        jtshape = "*?foo"
+        shape = (2, 3)
+    else:
+        jtshape = "?foo"
+        shape = (4,)
+
+    @jaxtyped
+    @typecheck
+    def f(
+        x: PyTree[Float[Array, jtshape], " T"], y: PyTree[Float[Array, jtshape], " T"]
+    ):
+        pass
+
+    x1 = jr.normal(getkey(), shape)
+    y1 = jr.normal(getkey(), shape)
+    x2 = jr.normal(getkey(), (5,))
+    y2 = jr.normal(getkey(), (5,))
+    f(x1, y1)
+    f((x1, x2), (y1, y2))
+
+    with pytest.raises(ParamError):
+        f(x1, y2)
+
+    with pytest.raises(ParamError):
+        f((x1, x2), (y2, y1))
+
+
+@pytest.mark.parametrize("variadic", (False, True))
+def test_treepath_dependence_dataclass(variadic, typecheck, getkey):
+    if variadic:
+        jtshape = "*?foo"
+        shape = (2, 3)
+    else:
+        jtshape = "?foo"
+        shape = (4,)
+
+    @jaxtyping._decorator._jaxtyped_typechecker(typecheck)
+    class A(eqx.Module):
+        x: PyTree[Float[Array, jtshape], " T"]
+        y: PyTree[Float[Array, jtshape], " T"]
+
+    x1 = jr.normal(getkey(), shape)
+    y1 = jr.normal(getkey(), shape)
+    x2 = jr.normal(getkey(), (5,))
+    y2 = jr.normal(getkey(), (5,))
+    A(x1, y1)
+    A((x1, x2), (y1, y2))
+
+    with pytest.raises(ParamError):
+        A(x1, y2)
+
+    with pytest.raises(ParamError):
+        A((x1, x2), (y2, y1))
+
+
+def test_treepath_dependence_missing_structure_annotation(typecheck, getkey):
+    @jaxtyped
+    @typecheck
+    def f(x: PyTree[Float[Array, "?foo"], " T"], y: PyTree[Float[Array, "?foo"]]):
+        pass
+
+    x1 = jr.normal(getkey(), (2,))
+    y1 = jr.normal(getkey(), (2,))
+    with pytest.raises(ValueError, match="except when contained with structured"):
+        f(x1, y1)
+
+
+def test_treepath_dependence_multiple_structure_annotation(typecheck, getkey):
+    @jaxtyped
+    @typecheck
+    def f(x: PyTree[PyTree[Float[Array, "?foo"], " S"], " T"]):
+        pass
+
+    x1 = jr.normal(getkey(), (2,))
+    with pytest.raises(ValueError, match="ambiguous which PyTree"):
+        f(x1)