Repository: patrick-kidger/jaxtyping (language: python)

Issue:
Type checking on closures
Hey @patrick-kidger,
great package as always. I was wondering what the best way is to type check closure functions? For example:

```python
import numpy as np
from jaxtyping import Float
from typeguard import typechecked


def f(n):
    @typechecked
    def g(x: Float[np.ndarray, f"dim1 {n}"]) -> Float[np.ndarray, f"{n}"]:
        return x.sum(0)

    return g


print(f(10)(np.random.randn(10, 9)))
```

This seems to work, but I'm wondering if this is the proper way? Or maybe something with `typing.TypeVar` is better?

I should also mention I'm not necessarily interested in closures, but more in class methods, where I want the return type to depend on a class parameter (e.g., `-> Float[Array, f"{self.latent}"]`).


Discussion / hints:
e.g.,

```python
import torch
from torch import nn, Tensor
import numpy as np
from jaxtyping import Float
from typeguard import typechecked

class Model(nn.Module):
    def __init__(self, latent):
        super().__init__()
        self.latent = latent
        self.net = nn.Linear(10, latent)

    @typechecked
    def forward(
        self, x: Float[Tensor, "batch 10"]
    ) -> Float[Tensor, f"batch {self.latent}"]:
        return self.net(x)

Model(5)(torch.randn(10, 10))
```

gives the error:

```
Traceback (most recent call last):
  File "/Users/mcranmer/Documents/perceiver4data/test.py", line 16, in <module>
    class Model(nn.Module):
  File "/Users/mcranmer/Documents/perceiver4data/test.py", line 24, in Model
    ) -> Float[Tensor, f"batch {self.latent}"]:
NameError: name 'self' is not defined
```
Hmm. So the closure example you've written looks correct. The class example errors because the annotation is evaluated when the function is created, not when the function is ran.

You could probably make this work by doing:
```python
@jaxtyped
@typechecked
def forward(self, x: Float[Tensor, "batch 10"]) -> Float[Tensor, "batch latent"]:
    assert isinstance(torch.zeros(self.latent), Float[Tensor, "latent"])
    return self.net(x)
```
to bind `latent` to the value of `self.latent`. (Each collection of named axes is valid for the dynamic scope of the `jaxtyped` decorator, and this includes the interior of the body function.)

It's obviously kind of hackish to have to create a tensor just for that purpose though. I'd be happy to take a PR offering a nicer syntax.
I have a similar piece of code that has the same issue but doesn't require a class:
```py
import torch
from jaxtyping import Float

def square(i: int) -> Float[torch.Tensor, "i"]:
    return torch.arange(i) ** 2
```

Maybe adding a function with the following type signature
```py
@overload
def assign_dims(**kwargs: int) -> None: ...

@overload
def assign_dims(names: Tuple[str], shape: Union[Tuple[int], int]) -> None: ...

@overload
def assign_dims(names: str, shape: Union[Tuple[int], int]) -> None: ...

def assign_dims(names, shape, /, **kwargs):
   ...

def square(i: int) -> Float[torch.Tensor, "i"]:
    assign_dims("i", i)
    return torch.arange(i) ** 2
```
Implementing this function is simple and straightforward. It requires gluing these two code snippets together:
https://github.com/google/jaxtyping/blob/a6ab6c0d28a5209e08c31791a1a20dc6cace6697/jaxtyping/_array_types.py#L177-L184
https://github.com/google/jaxtyping/blob/a6ab6c0d28a5209e08c31791a1a20dc6cace6697/jaxtyping/_array_types.py#L130-L137

Or a type annotation similar to the other ones in this repo would be useful, but I have no idea how to get static type checkers to be OK with this, but it does look more understandable to me:
```py
def square(i: Dims["i"]) -> Float[torch.Tensor, "i"]:
    return torch.arange(i) ** 2
```

@patrick-kidger Are we forced to use the function version or is there a clever way to get something similar to the type annotation version working?
I think we can probably convince static type checkers to at least ignore the annotation. (We might not be able to make them treat it as an `int` though.)

The real question is what kind of annotation we might dream up to support the use-case with classes, like in the original question. What syntax would be neatest?
Regardless of how the type annotation will be done, I am starting to think that a function for asserting the shape of an array (or a mechanism similar) is useful in its own right for introducing/asserting shape variables in the middle of a function, in addition to providing a way to provide shape variables for class attributes.

```py
def bincount_wrapper(x: Int[torch.Tensor, "elems"]) -> Int[torch.Tensor, "max_elem"]:
    print(x)
    assign_dims("max_elem", torch.amax(x))
    return torch.bincount(x)
```
Additional type annotations for class attributes could probably be implemented using the `assign_dims` function.
That sounds reasonable to me!
I am not entirely sure about the name of this function. Maybe changing its name to `assert_shape` or `assert_shapes` to mirror this function in TensorFlow would be a good idea? But it may also be confusing since the two functions are a little bit different.

https://www.tensorflow.org/api_docs/python/tf/debugging/assert_shapes

I am open for other name suggestions.
`bind_name` maybe? I don't have a strong opinion.
I like `bind_name`. I think implementing `bind_name` should be trivial. I can open a PR over the weekend.

Hi, 
I wannt to return back to discuss this solution: 


> Hmm. So the closure example you've written looks correct. The class example errors because the annotation is evaluated when the function is created, not when the function is ran.
> 
> You could probably make this work by doing:
> 
> ```python
> @jaxtyped
> @typechecked
> def forward(self, x: Float[Tensor, "batch 10"]) -> Float[Tensor, "batch latent"]:
>     assert isinstance(torch.zeros(self.latent), Float[Tensor, "latent"])
>     return self.net(x)
> ```
> 
> to bind `latent` to the value of `self.latent`. (Each collection of named axes is valid for the dynamic scope of the `jaxtyped` decorator, and this includes the interior of the body function.)
> 
> It's obviously kind of hackish to have to create a tensor just for that purpose though. I'd be happy to take a PR offering a nicer syntax.

we could maybe replace `isinstance` to functions like `assign_dims`,  but I wonder if it's possible to integreate this feature into @jaxtyped,  then it would be like again 

```
@jaxtyped
@typechecked
def forward(self, x: Float[Tensor, "batch 10"]) -> Float[Tensor, "batch self.latent"]:
    return self.net(x)
```

I think this would be more easy for reading and implementation. Otherwise if we manually bind `axis name`(i.e. `latent` here) and class member (i.e. `self.latent`) inside of function,  adding one sentence to assign dim is as same as simply doing as follows,   
```
assert output.shape[1] == self.latent
```  

so it would be really very appreciated to integrate accessing class member for type checking.





I didn't get around to writing test cases for this since I was busy recently. I can fish out the code I wrote for `assign_dims` sometime tomorrow.

I would hope that the `isinstance` method would work, but I would need to verify again. The issue I find with it is that it unnecessarily creates a tensor object that isn't used.

The problem I see with what your proposal is that the variables in the shape annotations are in a different scope than the variables in the function body. We could add `self` to the shape annotation scope, but it doesn't answer how to deal with variables from the function body or class variables in `classmethod`s.

The problem with `assert something.shape == myshape` (along with the different scopes) is that I think it would require expensive reparsing and inspection of the source code of the function to make work, but if it is possible to cheaply perform that operation, I would much prefer it.

Thank you for looking into this. I am not 100% pleased with the `bind_name` function myself, but I think it is the simplest and most general solution to this problem. I would like to be proven wrong. :)
Hi, @anivegesana 

Thanks for the fast reply. In my test, use `isinstance` works. But I don't think it's a good choice, for example, 
```
    @jaxtyped
    @typechecked
    def forward(self, x: Float[Tensor, "batch dim_in"], ) -> Float[Tensor, "batch latent"]:
        assert isinstance(torch.zeros(5), Float[Tensor, "latent"])
        assert isinstance(torch.zeros(4), Float[Tensor, "latent"])

        self.latent = 5
        return torch.zeros(8, 5)
```
In this case, at the second `assert`, it will throw error. If we swap two `assert` sentence, still the second throw error. What I want to say is that here `isinstance` works not like its name: checking if it's the instance of given class, it just serves as bind names with variable. And assert will never works(in first assert sentence). This could be confusing for new users. And I thinks it's not safe to use isinstance() directly. if we only annotate input and output, the relationship and constraint of symbols are easy to check and understand. But if justifications like "asser isinstance()" sparsely distributed inside function, it would be soonly confusing and difficult for debug. People has to be carefully check where it's initialized and where is the shape conflict. 

But I would like to directly add dim info to `memo`, as https://github.com/google/jaxtyping/issues/93#issuecomment-1622649793 also mentioned, in _array_types, it uses `single memo` and `variadic_memo`. In my understanding, the decorater `@jaxtyped`, it summarized anotation and create `memo`, then put the summerized infos into memo. (@patrick-kidger , do I understand right? or the mechanism is more complex?) It's more easy and clear to add two functions, `check dim` and `assign dim`, for parameters like input, all shape infos are determined. Check dim just compare "symbol" and given variable. For assigin dim, we only add symbol into `memo`, e.g.
```
DIM_0 = 10
def repeat_i_times(i: Int, x: Float[Tensor, "dim0"]) -> Float[Tensor, "dim1 dim0"]:
    check_dim("dim0", DIM_0) # assum x has shape same with global variable "DIM_0"
    assign_dim("dim1", i)
    return x.repeat(i, 1)   # output has shape (i, dim0)
```

Of course these two functions only make sense with `@jaxtyped`. 

And back to parse "self.xxx", Yes, it doesn't solve problem with variable inside function body and class method. In the latter two cases it could only be implemented by assign_name.  The motivation is that, for my experiences like machine learning, most of shape check are about checking like "batch data_dim", dimensions like "batch" are dynamic and flexible. Dimensions like "data_dim" are determined once the class is initailized.  So I think this feature could be a very helpful feature. 

The question is still there: what is the best way to annotate, so that it's most simple, robust and readable. I believe parsing "self.xxx", add function assign_dim() and check_dim() would be very useful. 


I think we are thinking the same thing. If there is no `@jaxtyped`, I chose to make `assign_dims` a no-op. That you described is literally how my `assign_dims` works. I attached the code and made a minor implementation to add a `check_dims`.

Treating `self` in the annotations is a little more involved and I will leave that for later. Since `assign_dims("xxx", self.xxx)` accomplishes this, I won't worry about it, but you can look into it if you like. I, personally, don't mind doing `assign_dims("batch num_channels height width", *image.shape)` in my code.
```python
from itertools import islice
from typing import SupportsInt, LiteralString

from jaxtyping._array_types import storage

def _dims(dims: LiteralString, /, *lengths: SupportsInt, assign: bool):
    no_temp_memo = hasattr(storage, "memo_stack") and len(storage.memo_stack) != 0
    if no_temp_memo:
        single_memo: dict[str, int]
        variadic_memo: dict[str, tuple[int, ...]]
        variadic_broadcast_memo: dict[str, list[tuple[int, ...]]]
        single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]

        dim_list = dims.split()
        iter_lengths = iter(lengths)
        seen_variadic = False

        for dim, length in zip(dim_list, iter_lengths, strict=True):
            length = int(length)

            if dim == "...":
                dim = "*_"

            if dim.startswith("*"):
                # We are dealing with a variadic
                if seen_variadic:
                    raise ValueError(f"Muliple variadic shape arguments in {dims!r}")
                seen_variadic = True

                amount_in_variadic = len(lengths) - len(dim_list)
                variadic = (length,) + tuple(int(i) for i in islice(iter_lengths, amount_in_variadic))

                if dim == "*_":
                    continue
                elif (old_variadic := single_memo.get(dim)) is not None:
                    if old_variadic != variadic:
                        raise ValueError(f"Conflict in dim {dim!r}. First found {old_variadic}. Now found {variadic}.")
                elif assign:
                    variadic_memo[dim] = variadic
            else:
                assert dim.isidentifier()

                if dim == "_":
                    continue
                elif (old_length := single_memo.get(dim)) is not None:
                    if old_length != length:
                        raise ValueError(f"Conflict in dim {dim!r}. First found {old_length}. Now found {length}.")
                elif assign:
                    single_memo[dim] = length

def assign_dims(dims: LiteralString, /, *lengths: SupportsInt):
    _dims(dims, *lengths, assign=True)

def check_dims(dims: LiteralString, /, *lengths: SupportsInt):
    _dims(dims, *lengths, assign=False)
```

It seems that I didn't add a case for variadic broadcasts. Am planning on looking into that later, but you can add support if you have the time.
In terms of making `self` work: I'm guessing this can probably be done by overriding `__get__`, which is part of Python's descriptor protocol.

In terms of the `assign_dims` functionality: I'd suggest that this actually make an `isinstance` call under the hood. And just use a dummy array type internal to jaxtyping:
```python
def assign_dim(dim: str, value: int):
    isinstance(DummyArray(value), Shaped[DummyArray, dim])
```
This avoids the needle to meddle with the internal machinery.
Yes. I was unaware that `isinstance` changes the state of the annotations. I had, at first, thought that it only checked the shape. I think we should go for that as well.
A small thing, for dummy array, I think using function like `np.empty(0, batch_size)` is better than `np.zeros(batchsize)`, since in this case  we don't need to consider memory cost and creating unused variable. 

Fix patch:
diff --git a/README.md b/README.md
index 3ec2f64..ebcbda8 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,7 @@ Type annotations **and runtime type-checking** for:
 ```python
 from jaxtyping import Array, Float, PyTree
 
-# Accepts floating-point 2D arrays with matching dimensions
+# Accepts floating-point 2D arrays with matching axes
 def matrix_multiply(x: Float[Array, "dim1 dim2"],
                     y: Float[Array, "dim2 dim3"]
                   ) -> Float[Array, "dim1 dim3"]:
diff --git a/docs/api/array.md b/docs/api/array.md
index ddbf109..0eadb12 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,12 +20,12 @@ When calling a function, variable-size axes and symbolic axes will be matched up
 
 In addition some modifiers can be applied:
 
-- Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `"*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.  
+- 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 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 _ _"`.
+- 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 _ _"`.
 - 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/).)
+- Prepend `?` to an axis to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree/).)
 
 When using multiple modifiers, their order does not matter.
 
@@ -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 dimensions:  
-    `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, "{size}"]:
+        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/docs/index.md b/docs/index.md
index 5511402..d7e59b8 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -24,7 +24,7 @@ The annotations provided by jaxtyping are compatible with runtime type-checking
 ```python
 from jaxtyping import Array, Float, PyTree
 
-# Accepts floating-point 2D arrays with matching dimensions
+# Accepts floating-point 2D arrays with matching axes
 def matrix_multiply(x: Float[Array, "dim1 dim2"],
                     y: Float[Array, "dim2 dim3"]
                   ) -> Float[Array, "dim1 dim3"]:
diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py
index 1f118c8..5ac236f 100644
--- a/jaxtyping/_array_types.py
+++ b/jaxtyping/_array_types.py
@@ -28,7 +28,12 @@
 import numpy as np
 
 from ._raise import jaxtyping_raise, jaxtyping_raise_from
-from ._storage import get_shape_memo, get_treepath_memo, set_shape_memo
+from ._storage import (
+    get_shape_memo,
+    get_treeflatten_memo,
+    get_treepath_memo,
+    set_shape_memo,
+)
 
 
 try:
@@ -89,10 +94,9 @@ def __init__(self, size, broadcastable):
 
 
 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 +114,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,15 +127,18 @@ 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 dimension '{cls_dim.elem_string}' 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 "
+                        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 "
                         "arguments."
                     ),
                     e,
@@ -153,32 +161,14 @@ def _check_dims(
     return True
 
 
-def _is_jax_extended_dtype(dtype: Any) -> bool:
-    if not has_jax:
-        return False
-    try:
-        is_dtype = issubclass(dtype, jax.numpy.generic)
-    except TypeError:
-        # `dtype` not a class
-        return False
-    else:
-        if is_dtype:
-            if hasattr(jax.dtypes, "extended"):  # jax>=0.4.14
-                return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
-            else:  # jax<=0.4.13
-                return jax.core.is_opaque_dtype(dtype)
-        else:
-            return False
-
-
 class _MetaAbstractArray(type):
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
+        if get_treeflatten_memo():
+            return True
 
-        if _is_jax_extended_dtype(obj.dtype):
-            dtype = str(obj.dtype)
-        elif hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
+        if hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
             # JAX, numpy
             dtype = obj.dtype.type.__name__
         elif hasattr(obj.dtype, "as_numpy_dtype"):
@@ -208,19 +198,24 @@ def __instancecheck__(cls, obj):
             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 +223,12 @@ def _check_shape(
         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 +236,10 @@ def _check_shape(
             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]
@@ -341,18 +337,18 @@ def _make_array(array_type, dim_str, dtypes, name):
         if "," in elem and "(" not in elem:
             # Common mistake.
             # Disable in the case that there's brackets to allow for function calls,
-            # e.g. `min(foo,bar)`, in symbolic dimensions.
-            raise ValueError("Dimensions should be separated with spaces, not commas")
+            # e.g. `min(foo,bar)`, in symbolic axes.
+            raise ValueError("Axes should be separated with spaces, not commas")
         if elem.endswith("#"):
             raise ValueError(
-                "As of jaxtyping v0.1.0, broadcastable dimensions are now denoted "
+                "As of jaxtyping v0.1.0, broadcastable axes 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; "
+                    "Anonymous multiple axes '...' must be used on its own; "
                     f"got {elem}"
                 )
             broadcastable = False
@@ -382,7 +378,7 @@ def _make_array(array_type, dim_str, dtypes, name):
                     if variadic:
                         raise ValueError(
                             "Do not use * twice to denote accepting multiple "
-                            "dimensions, e.g. `**foo` is not allowed"
+                            "axes, e.g. `**foo` is not allowed"
                         )
                     variadic = True
                     elem = elem[1:]
@@ -421,7 +417,7 @@ def _make_array(array_type, dim_str, dtypes, name):
         if variadic:
             if index_variadic is not None:
                 raise ValueError(
-                    "Cannot use multiple-dimension specifiers (`*name` or `...`) "
+                    "Cannot use variadic specifiers (`*name` or `...`) "
                     "more than once."
                 )
             index_variadic = index
@@ -429,7 +425,7 @@ def _make_array(array_type, dim_str, dtypes, name):
         if dim_type is _DimType.fixed:
             if variadic:
                 raise ValueError(
-                    "Cannot have a fixed axis bind to multiple dimensions, e.g. "
+                    "Cannot have a fixed axis bind to multiple axes, e.g. "
                     "`*4` is not allowed."
                 )
             if anonymous:
@@ -446,7 +442,7 @@ def _make_array(array_type, dim_str, dtypes, name):
             if anonymous:
                 if broadcastable:
                     raise ValueError(
-                        "Cannot have a dimension be both anonymous and "
+                        "Cannot have an axis be both anonymous and "
                         "broadcastable, e.g. `#_` is not allowed."
                     )
                 if variadic:
@@ -462,22 +458,20 @@ def _make_array(array_type, dim_str, dtypes, name):
             assert dim_type is _DimType.symbolic
             if anonymous:
                 raise ValueError(
-                    "Cannot have a symbolic dimension be anonymous, e.g. "
+                    "Cannot have a symbolic axis be anonymous, e.g. "
                     "`_foo+bar` is not allowed"
                 )
             if variadic:
                 raise ValueError(
-                    "Cannot have symbolic multiple-dimensions, e.g. "
+                    "Cannot have symbolic multiple-axes, e.g. "
                     "`*foo+bar` is not allowed"
                 )
             if treepath:
                 raise ValueError(
-                    "Cannot have a symbolic dimensions with tree-path dependence, e.g. "
+                    "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)
 
@@ -504,6 +498,16 @@ def _make_array(array_type, dim_str, dtypes, name):
             return array_type
         else:
             return _not_made
+    elif array_type is np.bool_:
+        if _check_scalar("bool", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is np.generic or array_type is np.number:
+        if _check_scalar("", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
     if issubclass(array_type, AbstractArray):
         if dtypes is _any_dtype:
             dtypes = array_type.dtypes
@@ -520,7 +524,7 @@ def _make_array(array_type, dim_str, dtypes, name):
                 index_variadic = array_type.index_variadic + len(dims)
             else:
                 raise ValueError(
-                    "Cannot use multiple-dimension specifiers (`*name` or `...`) "
+                    "Cannot use variadic specifiers (`*name` or `...`) "
                     "in both the original array and the extended array"
                 )
         dims = dims + array_type.dims
@@ -639,6 +643,7 @@ def __init_subclass__(cls, **kwargs):
         cls.dtypes = dtypes
 
 
+_prng_key = "prng_key"
 _bool = "bool"
 _bool_ = "bool_"
 _uint8 = "uint8"
@@ -707,8 +712,7 @@ class _Cls(AbstractDtype):
 Shaped = _make_dtype(_any_dtype, "Shaped")
 
 if has_jax:
-    _key_regex = re.compile(r"^key<\w+>$")
-    Key = _make_dtype(_key_regex, "Key")
+    Key = _make_dtype(_prng_key, "Key")
     # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
     # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype `uint32`.
     PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
index 8314cac..0bfa5c8 100644
--- a/jaxtyping/_decorator.py
+++ b/jaxtyping/_decorator.py
@@ -117,27 +117,6 @@ def f(x):
     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 f(...): ...
         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 __init__(self, *args, **kwargs):
             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 __init__(self, *args, **kwargs):
             # ```
             # 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 wrapped_fn(*args, **kwargs):  # pyright: ignore
             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.
@@ -305,13 +342,13 @@ def wrapped_fn(*args, **kwargs):
                                 name = fn.__name__
                             except AttributeError:
                                 name = fn.__class__.__name__
-                            paramstr = _remove_typing(param_signature)
+                            param_values = _pformat(bound.arguments, short_self=True)
+                            param_hints = _remove_typing(param_signature)
                             msg = (
                                 "Type-check error whilst checking the parameters of "
                                 f"{name}.{argmsg}\n"
-                                f"Called with args: {_pformat(args)}\n"
-                                f"Called with kwargs: {_pformat(kwargs)}\n"
-                                f"Parameter annotations: {paramstr}.\n"
+                                f"Called with arguments: {param_values}\n"
+                                f"Parameter annotations: {param_hints}.\n"
                                 + _exc_shape_info(memos)
                             )
                             raise TypeCheckError(msg) from e
@@ -345,23 +382,25 @@ def wrapped_fn(*args, **kwargs):
                                     name = fn.__name__
                                 except AttributeError:
                                     name = fn.__class__.__name__
-                                paramstr = _remove_typing(param_signature)
-                                returnstr = _remove_typing(
+                                param_values = _pformat(
+                                    bound.arguments, short_self=True
+                                )
+                                return_value = _pformat(out, short_self=False)
+                                param_hints = _remove_typing(param_signature)
+                                return_hint = _remove_typing(
                                     full_signature.return_annotation
                                 )
-                                if returnstr.startswith(
+                                if return_hint.startswith(
                                     "<class '"
-                                ) and returnstr.endswith("'>"):
-                                    returnstr = returnstr[8:-2]
-                                kwargs.pop(output_name)
+                                ) and return_hint.endswith("'>"):
+                                    return_hint = return_hint[8:-2]
                                 msg = (
                                     "Type-check error whilst checking the return value "
                                     f"of {name}.\n"
-                                    f"Called with args: {_pformat(args)}\n"
-                                    f"Called with kwargs: {_pformat(kwargs)}\n"
-                                    f"Return value: {_pformat(out)}\n"
-                                    f"Parameter annotations: {paramstr}.\n"
-                                    f"Return annotation: {returnstr}.\n"
+                                    f"Called with arguments: {param_values}\n"
+                                    f"Return value: {return_value}\n"
+                                    f"Parameter annotations: {param_hints}.\n"
+                                    f"Return annotation: {return_hint}.\n"
                                     + _exc_shape_info(memos)
                                 )
                                 raise TypeCheckError(msg) from e
@@ -376,7 +415,7 @@ def wrapped_fn(*args, **kwargs):
 
 class _JaxtypingContext:
     def __enter__(self):
-        push_shape_memo()
+        push_shape_memo({})
 
     def __exit__(self, exc_type, exc_value, exc_tb):
         pop_shape_memo()
@@ -388,7 +427,7 @@ def _check_dataclass_annotations(self, typechecker):
     `self` should be a dataclass instancae. `typechecker` should be e.g.
     `beartype.beartype` or `typeguard.typechecked`.
     """
-    parameters = []
+    parameters = [inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)]
     values = {}
     for field in dataclasses.fields(self):
         annotation = field.type
@@ -424,7 +463,7 @@ def _check_dataclass_annotations(self, typechecker):
         self.__class__.__name__, signature, module, output=False
     )
     f = jaxtyped(f, typechecker=typechecker)
-    f(**values)
+    f(self, **values)
 
 
 def _make_fn_with_signature(
@@ -618,7 +657,7 @@ def _remove_typing(x):
     return x
 
 
-def _pformat(x):
+def _pformat(x, short_self: bool):
     # No performance concerns from delayed imports -- this is only used when we're about
     # to raise an error anyway.
     try:
@@ -628,6 +667,14 @@ def _pformat(x):
         import equinox as eqx
 
         pformat = eqx.tree_pformat
+        if short_self:
+            try:
+                self = x["self"]
+            except KeyError:
+                pass
+            else:
+                is_self = lambda y: y is self
+                pformat = ft.partial(pformat, truncate_leaf=is_self)
     except Exception:
         import pprint
 
@@ -639,18 +686,26 @@ 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
+    single_memo = {
+        name: size
+        for name, size in single_memo.items()
+        if not name.startswith("~~delete~~")
+    }
+    variadic_memo = {
+        name: shape
+        for name, (_, shape) in variadic_memo.items()
+        if not name.startswith("~~delete~~")
+    }
     pieces = []
     if len(single_memo) > 0 or len(variadic_memo) > 0:
         pieces.append(
             "The current values for each jaxtyping axis annotation are as follows."
         )
         for name, size in single_memo.items():
-            if not name.startswith("~~delete~~"):
-                pieces.append(f"{name}={size}")
-        for name, (_, shape) in variadic_memo.items():
-            if not name.startswith("~~delete~~"):
-                pieces.append(f"{name}={shape}")
+            pieces.append(f"{name}={size}")
+        for name, shape in variadic_memo.items():
+            pieces.append(f"{name}={shape}")
     if len(pytree_memo) > 0:
         pieces.append(
             "The current values for each jaxtyping PyTree structure annotation are as "
diff --git a/jaxtyping/_import_hook.py b/jaxtyping/_import_hook.py
index f2a551e..10f7f69 100644
--- a/jaxtyping/_import_hook.py
+++ b/jaxtyping/_import_hook.py
@@ -109,7 +109,7 @@ def __init__(self, typechecker):
 
         elif typechecker is None:
             # If it is None, ignore it silently (use dummy decorator)
-            self.hash = 0
+            self.hash = "0"
             Typechecker.lookup[self.hash] = lambda x, *_, **__: x
         else:
             # Passed typechecker is invalid
diff --git a/jaxtyping/_pytree_type.py b/jaxtyping/_pytree_type.py
index 84ad9ec..df89997 100644
--- a/jaxtyping/_pytree_type.py
+++ b/jaxtyping/_pytree_type.py
@@ -26,9 +26,11 @@
 
 from ._raise import jaxtyping_raise_from
 from ._storage import (
+    clear_treeflatten_memo,
     clear_treepath_memo,
     get_shape_memo,
     set_shape_memo,
+    set_treeflatten_memo,
     set_treepath_memo,
 )
 
@@ -52,20 +54,29 @@ def __call__(self, *args, **kwargs):
     def __instancecheck__(cls, obj):
         if not hasattr(cls, "leaftype"):
             return True  # Just `isinstance(x, PyTree)`
+        # Handle beartype doing `isinstance(None, hint)` to check if
+        # is `instance`able.
+        if obj is None:
+            return True
 
-        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):
@@ -74,7 +85,7 @@ def _check(cls, obj, pytree_memo):
             def is_flatten_leaftype(x):
                 return False
 
-            def is_check_leaftype(x, new_scope):
+            def is_check_leaftype(x):
                 return True
 
         else:
@@ -88,22 +99,21 @@ def is_check_leaftype(x, new_scope):
             def accepts_leaftype(x: cls.leaftype):
                 pass
 
-            def is_leaftype(x, new_scope=True):
-                if new_scope and cls.structure is not None:
-                    set_treepath_memo(None, cls.structure)
+            def is_leaftype(x):
                 try:
                     accepts_leaftype(x)
                 except _TypeCheckError:
                     return False
                 else:
                     return True
-                finally:
-                    if new_scope and cls.structure is not None:
-                        clear_treepath_memo()
 
             is_flatten_leaftype = is_check_leaftype = is_leaftype
 
-        leaves, structure = jtu.tree_flatten(obj, is_leaf=is_flatten_leaftype)
+        set_treeflatten_memo()
+        try:
+            leaves, structure = jtu.tree_flatten(obj, is_leaf=is_flatten_leaftype)
+        finally:
+            clear_treeflatten_memo()
         if cls.structure is not None:
             if cls.structure.isidentifier():
                 try:
@@ -168,7 +178,7 @@ def is_leaftype(x, new_scope=True):
             for leaf_index, leaf in enumerate(leaves):
                 if cls.structure is not None:
                     set_treepath_memo(leaf_index, cls.structure)
-                if not is_check_leaftype(leaf, new_scope=False):
+                if not is_check_leaftype(leaf):
                     return False
                 clear_treepath_memo()
         finally:
diff --git a/jaxtyping/_raise.py b/jaxtyping/_raise.py
index 9230877..38d3f1a 100644
--- a/jaxtyping/_raise.py
+++ b/jaxtyping/_raise.py
@@ -6,6 +6,7 @@ def jaxtyping_raise(e) -> NoReturn:
     `TypeCheckError`. All `raise` statements from within `__instancecheck__` should use
     this.
     """
+    __tracebackhide__ = True
     try:
         raise e
     except Exception as f:
diff --git a/jaxtyping/_storage.py b/jaxtyping/_storage.py
index ed9b3bc..73d1f64 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
 
@@ -96,3 +104,21 @@ def get_treepath_memo() -> str:
             )
         )
     return _treepath_storage.value
+
+
+_treeflatten_storage = threading.local()
+
+
+def clear_treeflatten_memo() -> None:
+    _treeflatten_storage.value = False
+
+
+def set_treeflatten_memo():
+    _treeflatten_storage.value = True
+
+
+def get_treeflatten_memo():
+    try:
+        return _treeflatten_storage.value
+    except AttributeError:
+        return False


Test patch:
diff --git a/test/test_array.py b/test/test_array.py
index 4f3fc26..1b81cae 100644
--- a/test/test_array.py
+++ b/test/test_array.py
@@ -17,6 +17,7 @@
 # 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 dataclasses as dc
 import sys
 from typing import get_args, get_origin, Union
 
@@ -33,7 +34,10 @@
     Bool,
     Float,
     Float32,
+    jaxtyped,
+    Key,
     PRNGKeyArray,
+    Scalar,
     Shaped,
 )
 
@@ -448,6 +452,54 @@ def foo(x: Float32[Array, " 2*dim"]):
         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_deferred_symbolic_dataclass(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    @dc.dataclass
+    class A:
+        value: int
+        array: Float[Array, " {self.value}"]
+
+    A(3, jnp.zeros(3))
+
+    with pytest.raises(ParamError):
+        A(3, jnp.zeros(4))
+
+
 def test_arraylike(typecheck, getkey):
     floatlike1 = Float32[ArrayLike, ""]
     floatlike2 = Float[ArrayLike, ""]
@@ -459,22 +511,18 @@ def test_arraylike(typecheck, getkey):
     assert set(get_args(floatlike1)) == {
         Float32[Array, ""],
         Float32[np.ndarray, ""],
-        Float32[np.bool_, ""],
         Float32[np.number, ""],
         float,
     }
     assert set(get_args(floatlike2)) == {
         Float[Array, ""],
         Float[np.ndarray, ""],
-        Float[np.bool_, ""],
         Float[np.number, ""],
         float,
     }
     assert set(get_args(floatlike3)) == {
         Float32[Array, "4"],
         Float32[np.ndarray, "4"],
-        Float32[np.bool_, "4"],
-        Float32[np.number, "4"],
     }
 
     shaped1 = Shaped[ArrayLike, ""]
@@ -494,8 +542,6 @@ def test_arraylike(typecheck, getkey):
     assert set(get_args(shaped2)) == {
         Shaped[Array, "4"],
         Shaped[np.ndarray, "4"],
-        Shaped[np.bool_, "4"],
-        Shaped[np.number, "4"],
     }
 
 
@@ -551,8 +597,8 @@ def test_key(jaxtyp, typecheck):
     def f(x: PRNGKeyArray):
         pass
 
-    x = jr.PRNGKey(0)
-    f(x)
+    f(jr.key(0))
+    f(jr.PRNGKey(0))
 
     with pytest.raises(ParamError):
         f(object())
@@ -564,6 +610,30 @@ def f(x: PRNGKeyArray):
         f(jnp.array(3.0))
 
 
+def test_key_dtype(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f1(x: Key[Array, ""]):
+        pass
+
+    @jaxtyp(typecheck)
+    def f2(x: Key[Scalar, ""]):
+        pass
+
+    for f in (f1, f2):
+        f(jr.key(0))
+
+        with pytest.raises(ParamError):
+            f(jr.PRNGKey(0))
+        with pytest.raises(ParamError):
+            f(object())
+        with pytest.raises(ParamError):
+            f(1)
+        with pytest.raises(ParamError):
+            f(jnp.array(3))
+        with pytest.raises(ParamError):
+            f(jnp.array(3.0))
+
+
 def test_extension(jaxtyp, typecheck, getkey):
     X = Shaped[Array, "a b"]
     Y = Shaped[X, "c d"]
diff --git a/test/test_decorator.py b/test/test_decorator.py
index 331a3d5..52c0c9a 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):
@@ -108,7 +110,7 @@ def f(**kwargs):
 
 def test_defaults(jaxtyp, typecheck):
     @jaxtyp(typecheck)
-    def f(x, y=1):
+    def f(x: int, y=1):
         pass
 
     f(1)
diff --git a/test/test_messages.py b/test/test_messages.py
index 2d8f0ad..2568baf 100644
--- a/test/test_messages.py
+++ b/test/test_messages.py
@@ -15,8 +15,7 @@ def f(x: str, y: str, z: int):
     matches = [
         "Type-check error whilst checking the parameters of f",
         "The problem arose whilst typechecking argument 'z'.",
-        r"Called with args: \('hi', 'bye', 'not-an-int'\)",
-        "Called with kwargs: {}",
+        "Called with arguments: {'x': 'hi', 'y': 'bye', 'z': 'not-an-int'}",
         r"Parameter annotations: \(x: str, y: str, z: int\).",
     ]
     for match in matches:
@@ -32,8 +31,7 @@ def g(x: Float[Array, "a b"], y: Float[Array, "b c"]):
     matches = [
         "Type-check error whilst checking the parameters of g",
         "The problem arose whilst typechecking argument 'y'.",
-        r"Called with args: \(f32\[2,3\],\)",
-        r"Called with kwargs: {'y': f32\[4,3\]}",
+        r"Called with arguments: {'x': f32\[2,3\], 'y': f32\[4,3\]}",
         (
             r"Parameter annotations: \(x: Float\[Array, 'a b'\], y: "
             r"Float\[Array, 'b c'\]\)."
@@ -56,8 +54,7 @@ def f(x: PyTree[Any, " T"], y: PyTree[Any, " S"]) -> PyTree[Any, "T S"]:
     y = {"a": 1}
     matches = [
         "Type-check error whilst checking the return value of f",
-        r"Called with args: \(\(1, 2\),\)",
-        r"Called with kwargs: {'y': {'a': 1}}",
+        r"Called with arguments: {'x': \(1, 2\), 'y': {'a': 1}}",
         "Return value: 'foo'",
         r"Return annotation: PyTree\[Any, \"T S\"\].",
         (
@@ -86,13 +83,12 @@ class M(eqx.Module):
     matches = [
         "Type-check error whilst checking the parameters of M",
         "The problem arose whilst typechecking argument 'z'.",
-        r"Called with args: \(\)",
         (
-            r"Called with kwargs: {'x': f32\[2,3\], 'y': \(1, \(3, 4\)\), "
-            r"'z': 'not-an-int'}"
+            r"Called with arguments: {'self': M\(\.\.\.\), 'x': f32\[2,3\], "
+            r"'y': \(1, \(3, 4\)\), 'z': 'not-an-int'}"
         ),
         (
-            r"Parameter annotations: \(x: Float\[Array, '\*foo'\], "
+            r"Parameter annotations: \(self: Any, x: Float\[Array, '\*foo'\], "
             r"y: PyTree\[Any, \"T\"\], z: int\)."
         ),
         "The current values for each jaxtyping axis annotation are as follows.",