Symbolic expressions now support delayed binding to arguments. Fixes #93.

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
 

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, "<string>", "eval")
-            elem = _SymbolicDim(elem, broadcastable, elem_string)
+            elem = _SymbolicDim(elem, broadcastable)
         dims.append(elem)
     dims = tuple(dims)
 

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(

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):

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
 

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

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):