diff --git a/docs/api/array.md b/docs/api/array.md index 8a38b69..d491944 100644 --- a/docs/api/array.md +++ b/docs/api/array.md @@ -10,8 +10,9 @@ The shape should be a string of space-separated symbols, such as `"a b c d"`. Ea - `int`: fixed-size axis, e.g. `"28 28"`. - `str`: variable-size axis, e.g. `"channels"`. -- A symbolic expression (without spaces!) in terms of other variable-size axes, e.g. - `def remove_last(x: Float[Array, "dim"]) -> Float[Array, "dim-1"]`. +- A symbolic expression in terms of other variable-size axes, e.g. + `def remove_last(x: Float[Array, "dim"]) -> Float[Array, "dim-1"]`. + Symbolic expressions must not use any spaces, otherwise each piece is treated as as a separate axis. 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.md).) @@ -19,7 +20,7 @@ When calling a function, variable-size axes and symbolic axes will be matched up In addition some modifiers can be applied: -- Prepend `*` to an axis to indicate that it can match multiple axes, e.g. `"*batch c h w"` will match zero or more batch axes. +- Prepend `*` to an axis to indicate that it can match multiple axes, e.g. `"*batch"` will match zero or more batch axes. - Prepend `#` to an axis to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g. `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`. - Prepend `_` to an axis to disable any runtime checking of that axis (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`. @@ -37,9 +38,18 @@ As a special case: - To denote a scalar shape use `""`, e.g. `Float[Array, ""]`. - To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `Float[Array, "..."]`. - 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 axes: - `def add(x: Float[Array, "*#foo"], y: Float[Array, "*#foo"]) -> Float[Array, "*#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. +- Symbolic expressions are evaluated in two stages: they are first evaluated as f-strings using the arguments of the function, and second are evaluated using the processed axis sizes. The f-string evaluation means that they can use local variables by enclosing them with curly braces, e.g. `{variable}`, e.g. + ```python + def full(size: int, fill: float) -> Float[Array, "{shape}"]: + return jax.numpy.full((size,), fill) + + class SomeClass: + some_value = 5 + + def full(self, fill: float) -> Float[Array, "{self.some_value}+3"]: + return jax.numpy.full((self.some_value + 3,), fill) + ``` ## Dtype diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py index 64f95ca..9214014 100644 --- a/jaxtyping/_array_types.py +++ b/jaxtyping/_array_types.py @@ -89,10 +89,9 @@ class _FixedDim: class _SymbolicDim: - def __init__(self, expr, broadcastable, elem_string): - self.expr = expr + def __init__(self, elem, broadcastable): + self.elem = elem self.broadcastable = broadcastable - self.elem_string = elem_string _AbstractDimOrVariadicDim = Union[ @@ -110,6 +109,7 @@ def _check_dims( cls_dims: list[_AbstractDim], obj_shape: tuple[int, ...], single_memo: dict[str, int], + arg_memo: dict[str, Any], ) -> bool: assert len(cls_dims) == len(obj_shape) for cls_dim, obj_size in zip(cls_dims, obj_shape): @@ -122,12 +122,15 @@ def _check_dims( return False elif type(cls_dim) is _SymbolicDim: try: + # Support f-string syntax. + # https://stackoverflow.com/a/53671539/22545467 + elem = eval(f"f'{cls_dim.elem}'", arg_memo.copy()) # Make a copy to avoid `__builtins__` getting added as a key. - eval_size = eval(cls_dim.expr, single_memo.copy()) + eval_size = eval(elem, single_memo.copy()) except NameError as e: jaxtyping_raise_from( NameError( - f"Cannot process symbolic axis '{cls_dim.elem_string}' as " + f"Cannot process symbolic axis '{cls_dim.elem}' as " "some axis names have not been processed. In practice you " "should usually only use symbolic axes in annotations " "for return types, referring only to axes annotated for " @@ -208,19 +211,24 @@ class _MetaAbstractArray(type): if not in_dtypes: return False - single_memo, variadic_memo, pytree_memo = get_shape_memo() + single_memo, variadic_memo, pytree_memo, arg_memo = get_shape_memo() single_memo_bak = single_memo.copy() variadic_memo_bak = variadic_memo.copy() pytree_memo_bak = pytree_memo.copy() + arg_memo_bak = arg_memo.copy() try: - check = cls._check_shape(obj, single_memo, variadic_memo) + check = cls._check_shape(obj, single_memo, variadic_memo, arg_memo) except Exception: - set_shape_memo(single_memo_bak, variadic_memo_bak, pytree_memo_bak) + set_shape_memo( + single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak + ) raise if check: return True else: - set_shape_memo(single_memo_bak, variadic_memo_bak, pytree_memo_bak) + set_shape_memo( + single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak + ) return False def _check_shape( @@ -228,11 +236,12 @@ class _MetaAbstractArray(type): obj, single_memo: dict[str, int], variadic_memo: dict[str, tuple[bool, tuple[int, ...]]], + arg_memo: dict[str, Any], ): if cls.index_variadic is None: if obj.ndim != len(cls.dims): return False - return _check_dims(cls.dims, obj.shape, single_memo) + return _check_dims(cls.dims, obj.shape, single_memo, arg_memo) else: if obj.ndim < len(cls.dims) - 1: return False @@ -240,10 +249,10 @@ class _MetaAbstractArray(type): j = -(len(cls.dims) - i - 1) if j == 0: j = None - if not _check_dims(cls.dims[:i], obj.shape[:i], single_memo): + if not _check_dims(cls.dims[:i], obj.shape[:i], single_memo, arg_memo): return False if j is not None and not _check_dims( - cls.dims[j:], obj.shape[j:], single_memo + cls.dims[j:], obj.shape[j:], single_memo, arg_memo ): return False variadic_dim = cls.dims[i] @@ -475,9 +484,7 @@ def _make_array(array_type, dim_str, dtypes, name): "Cannot have a symbolic axis with tree-path dependence, e.g. " "`?foo+bar` is not allowed" ) - elem_string = elem - elem = compile(elem, "", "eval") - elem = _SymbolicDim(elem, broadcastable, elem_string) + elem = _SymbolicDim(elem, broadcastable) dims.append(elem) dims = tuple(dims) diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index 8314cac..dd57a72 100644 --- a/jaxtyping/_decorator.py +++ b/jaxtyping/_decorator.py @@ -117,27 +117,6 @@ def jaxtyped(fn=None, *, typechecker=None): If `fn` is a dataclass, then `fn` is returned directly, and additionally its `__init__` method is wrapped and modified in-place. - **Notes for advanced users** - - Put precisely, the axis names in e.g. `Float[Array, "batch channels"]` and the - structure names in e.g. `PyTree[int, "T"]` are all scoped to the thread-local - dynamic context of a `jaxtyped`-wrapped function. A new dynamic context will allow - different values to be bound to the same name. After this new dynamic context is - finished then the old one is returned to. - - Binding of a value against a name is done with an `isinstance` check, for example - `isinstance(jnp.zeros((3, 4)), Float[Array, "dim1 dim2"])` will bind `dim1=3` and - `dim2=4`. - - This means you can use `isinstance` checks inside a function body and have them - contribute to the same collection of consistency checks performed by a typechecker - against its arguments. (Or even forgo a typechecker that analyses arguments, and - instead just do your own manual `isinstance` checks.) - - Only `isinstance` checks that pass will contribute to the store of values; those - that fail will not. As such it is safe to write e.g. - `assert not isinstance(x, Float32[Array, "foo"])`. - !!! Info "Old syntax" jaxtyping previously (before v0.2.24) recommended using this double-decorator @@ -151,6 +130,53 @@ def jaxtyped(fn=None, *, typechecker=None): discussed above will produce easier-to-debug error messages. Under the hood, the new syntax more carefully manipulates the typechecker so as to determine where a type-check error arises. + + ??? Info "Notes for advanced users" + + **Dynamic contexts:** + + Put precisely, the axis names in e.g. `Float[Array, "batch channels"]` and the + structure names in e.g. `PyTree[int, "T"]` are all scoped to the thread-local + dynamic context of a `jaxtyped`-wrapped function. If from within that function + we then call another `jaxtyped`-wrapped function, then a new context is pushed + to the stack. The axis sizes and PyTree structures of this inner function will + then not be compared against the axis sizes and PyTree structures of the outer + function. After the inner function returns then this inner context is popped + from the stack, and the previous context is returned to. + + **isinstance:** + + Binding of a value against a name is done with an `isinstance` check, for + example `isinstance(jnp.zeros((3, 4)), Float[Array, "dim1 dim2"])` will bind + `dim1=3` and `dim2=4`. In practice these `isinstance` checks are usually done by + the run-time typechecker `typechecker` that is supplied as an argument. + + This can also be done manually: add `isinstance` checks inside a function body + and they will contribute to the same collection of consistency checks as are + performed by the typechecker on the arguments and return values. (Or you can + forgo such a typechecker altogether -- i.e. `typechecker=None` -- and only do + your own manual `isinstance` checks.) + + Only `isinstance` checks that pass will contribute to the store of values; those + that fail will not. As such it is safe to write e.g. + `assert not isinstance(x, Float32[Array, "foo"])`. + + **Decoupling contexts from function calls:** + + If you would like a new dynamic context *without* calling a new function, then + `jaxtyped` may be passed the string `"context"` and used as a context manager: + ```python + with jaxtyped("context"): + assert isinstance(x, Float[Array, "batch channel"]) + ``` + which is equivalent to placing this code inside a new function wrapped in + `jaxtyped(typechecker=None)`. Usage like this is very rare; it's mostly only + useful when working at the global scope. + + Conversely, if you would like to call a new function *without* creating a new + dynamic context (and using the same set of axis and structure values), then + simply do not add a `jaxtyped` decorator to your inner function, whilst + continuing to perform type-checking in whatever way you prefer. """ if fn is None: @@ -210,6 +236,13 @@ def jaxtyped(fn=None, *, typechecker=None): fdel = jaxtyped(fn.fdel, typechecker=typechecker) return property(fget=fget, fset=fset, fdel=fdel) elif fn == "context": + if typechecker is not None: + raise ValueError( + "Cannot use `jaxtyped` as a context with a typechecker. That is, " + "`with jaxtyped('context', typechecker=...):`. is not allowed. In this " + "case the type checker does not actually do anything, as there is no " + "function to type-check." + ) return _JaxtypingContext() else: if typechecker is None: @@ -221,9 +254,13 @@ def jaxtyped(fn=None, *, typechecker=None): # ``` # in which case make a best-effort attempt to add shape information for any # type errors. + + signature = inspect.signature(fn) + @ft.wraps(fn) def wrapped_fn(*args, **kwargs): # pyright: ignore - memos = push_shape_memo() + bound = signature.bind(*args, **kwargs) + memos = push_shape_memo(bound.arguments) try: return fn(*args, **kwargs) except Exception as e: @@ -286,9 +323,9 @@ def jaxtyped(fn=None, *, typechecker=None): def wrapped_fn(*args, **kwargs): # Raise bind-time errors before we do any shape analysis. (I.e. skip # the pointless jaxtyping information for a non-typechecking failure.) - param_signature.bind(*args, **kwargs) + bound = param_signature.bind(*args, **kwargs) - memos = push_shape_memo() + memos = push_shape_memo(bound.arguments) try: # First type-check just the parameters before the function is # called. @@ -376,7 +413,7 @@ def jaxtyped(fn=None, *, typechecker=None): class _JaxtypingContext: def __enter__(self): - push_shape_memo() + push_shape_memo({}) def __exit__(self, exc_type, exc_value, exc_tb): pop_shape_memo() @@ -639,7 +676,7 @@ def _exc_shape_info(memos) -> str: """Gives debug information on the current state of jaxtyping's internal memos. Used in type-checking error messages. """ - single_memo, variadic_memo, pytree_memo = memos + single_memo, variadic_memo, pytree_memo, _ = memos pieces = [] if len(single_memo) > 0 or len(variadic_memo) > 0: pieces.append( diff --git a/jaxtyping/_pytree_type.py b/jaxtyping/_pytree_type.py index 84ad9ec..5526296 100644 --- a/jaxtyping/_pytree_type.py +++ b/jaxtyping/_pytree_type.py @@ -53,19 +53,24 @@ 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, variadic_memo, pytree_memo, arg_memo = get_shape_memo() single_memo_bak = single_memo.copy() variadic_memo_bak = variadic_memo.copy() pytree_memo_bak = pytree_memo.copy() + arg_memo_bak = arg_memo.copy() try: out = cls._check(obj, pytree_memo) except Exception: - set_shape_memo(single_memo_bak, variadic_memo_bak, pytree_memo_bak) + set_shape_memo( + single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak + ) raise if out: return True else: - set_shape_memo(single_memo_bak, variadic_memo_bak, pytree_memo_bak) + set_shape_memo( + single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak + ) return False def _check(cls, obj, pytree_memo): diff --git a/jaxtyping/_storage.py b/jaxtyping/_storage.py index ed9b3bc..84354f6 100644 --- a/jaxtyping/_storage.py +++ b/jaxtyping/_storage.py @@ -18,7 +18,7 @@ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. import threading -from typing import Optional +from typing import Any, Optional from ._raise import jaxtyping_raise @@ -32,7 +32,9 @@ def _has_shape_memo(): def get_shape_memo(): if _has_shape_memo(): - single_memo, variadic_memo, pytree_memo = _shape_storage.memo_stack[-1] + single_memo, variadic_memo, pytree_memo, arguments = _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 @@ -40,21 +42,27 @@ def get_shape_memo(): single_memo = {} variadic_memo = {} pytree_memo = {} - return single_memo, variadic_memo, pytree_memo + arguments = {} + return single_memo, variadic_memo, pytree_memo, arguments -def set_shape_memo(single_memo, variadic_memo, pytree_memo) -> None: +def set_shape_memo(single_memo, variadic_memo, pytree_memo, arg_memo) -> None: if _has_shape_memo(): - _shape_storage.memo_stack[-1] = single_memo, variadic_memo, pytree_memo + _shape_storage.memo_stack[-1] = ( + single_memo, + variadic_memo, + pytree_memo, + arg_memo, + ) -def push_shape_memo(): +def push_shape_memo(arguments: dict[str, Any]): 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 = [] - memos = ({}, {}, {}) + memos = ({}, {}, {}, arguments.copy()) memo_stack.append(memos) return memos diff --git a/test/test_array.py b/test/test_array.py index 4f3fc26..52b6e99 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -448,6 +448,41 @@ def test_incomplete_symbolic(jaxtyp, typecheck, getkey): foo(x) +def test_deferred_symbolic_good(jaxtyp, typecheck): + @jaxtyp(typecheck) + def foo(dim: int, fill: Float[Array, ""]) -> Float[Array, " {dim}"]: + return jnp.full((dim,), fill) + + class A: + size = 5 + + @jaxtyp(typecheck) + def bar(self, fill: Float[Array, ""]) -> Float[Array, " {self.size}"]: + return jnp.full((self.size,), fill) + + foo(3, jnp.array(0.0)) + A().bar(jnp.array(0.0)) + + +def test_deferred_symbolic_bad(jaxtyp, typecheck): + @jaxtyp(typecheck) + def foo(dim: int, fill: Float[Array, ""]) -> Float[Array, " {dim-1}"]: + return jnp.full((dim,), fill) + + class A: + size = 5 + + @jaxtyp(typecheck) + def bar(self, fill: Float[Array, ""]) -> Float[Array, " {self.size}-1"]: + return jnp.full((self.size,), fill) + + with pytest.raises(ReturnError): + foo(3, jnp.array(0.0)) + + with pytest.raises(ReturnError): + A().bar(jnp.array(0.0)) + + def test_arraylike(typecheck, getkey): floatlike1 = Float32[ArrayLike, ""] floatlike2 = Float[ArrayLike, ""] diff --git a/test/test_decorator.py b/test/test_decorator.py index 331a3d5..3d0001f 100644 --- a/test/test_decorator.py +++ b/test/test_decorator.py @@ -88,6 +88,8 @@ def test_context(getkey): with jaxtyped("context"): assert isinstance(a, Float[Array, "foo bar"]) assert not isinstance(b, Float[Array, "foo"]) + assert isinstance(a, Float[Array, "foo bar"]) + assert isinstance(b, Float[Array, "foo"]) def test_varargs(jaxtyp, typecheck):