From d12291de7e4f04a750f0e92e029486a7e4c87c6f Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Sat, 21 Oct 2023 19:53:54 -0700 Subject: [PATCH] Added support for treepath-dependent sizes. --- docs/api/array.md | 1 + docs/api/pytree.md | 32 +++++++++++++++ jaxtyping/_array_types.py | 68 +++++++++++++++++++++++-------- jaxtyping/_decorator.py | 45 +++------------------ jaxtyping/_pytree_type.py | 45 +++++++++++++++++---- jaxtyping/_storage.py | 85 +++++++++++++++++++++++++++++++++++++++ test/test_pytree.py | 84 +++++++++++++++++++++++++++++++++++++- 7 files changed, 296 insertions(+), 64 deletions(-) create mode 100644 jaxtyping/_storage.py diff --git a/docs/api/array.md b/docs/api/array.md index d64aa48..fda2dbd 100644 --- a/docs/api/array.md +++ b/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. diff --git a/docs/api/pytree.md b/docs/api/pytree.md index 2933fa5..fa864f0 100644 --- a/docs/api/pytree.md +++ b/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. diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py index 519e0c5..22e48f1 100644 --- a/jaxtyping/_array_types.py +++ b/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, "", "eval") elem = _SymbolicDim(elem, broadcastable, elem_string) diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index 9b71c82..1d907f6 100644 --- a/jaxtyping/_decorator.py +++ b/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 diff --git a/jaxtyping/_pytree_type.py b/jaxtyping/_pytree_type.py index 336227a..e0a4aec 100644 --- a/jaxtyping/_pytree_type.py +++ b/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 diff --git a/jaxtyping/_storage.py b/jaxtyping/_storage.py new file mode 100644 index 0000000..63cac07 --- /dev/null +++ b/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 diff --git a/test/test_pytree.py b/test/test_pytree.py index c1459f5..169f05c 100644 --- a/test/test_pytree.py +++ b/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)