Typecheck errors now state the size of the stored axis and structure values.

jaxtyping/_decorator.py

@@ -20,7 +20,7 @@
 import dataclasses
 import functools as ft
 import inspect
-import textwrap
+import sys
 import types
 import weakref
 from typing import get_args, get_origin
@@ -34,7 +34,7 @@ else:
     traceback_util.register_exclusion(__file__)
 
 
-from ._storage import pop_shape_memo, push_shape_memo
+from ._storage import get_shape_memo, pop_shape_memo, push_shape_memo
 
 
 _jaxtyped_fns = weakref.WeakSet()
@@ -88,6 +88,8 @@ def jaxtyped(fn):
         return fn
     elif inspect.isclass(fn):  # allow decorators on class definitions
         if dataclasses.is_dataclass(fn):
+            # TODO(kidger): unify this branch with `_jaxtyped_typechecker` below,
+            # perhaps if we ever do typechecking on an argument-by-argument basis.
             init = jaxtyped(fn.__init__)
             fn.__init__ = init
             return fn
@@ -123,9 +125,20 @@ def jaxtyped(fn):
 
         @ft.wraps(fn)
         def wrapped_fn(*args, **kwargs):
-            push_shape_memo()
+            memos = push_shape_memo()
             try:
                 return fn(*args, **kwargs)
+            except Exception as e:
+                if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
+                    shape_info = _exc_shape_info(memos)
+                    if shape_info != "":
+                        msg = (
+                            "The preceding error occurred within the scope of a "
+                            "`jaxtyping.jaxtyped` function, and may be due to a "
+                            "typecheck error. "
+                        )
+                        e.add_note(_jaxtyping_note_str(_spacer + msg + shape_info))
+                raise
             finally:
                 pop_shape_memo()
 
@@ -173,18 +186,24 @@ def _check_dataclass_annotations(self, typechecker):
         except AttributeError:
             continue  # allow uninitialised fields, which are allowed on dataclasses
 
-        # Dynamic `exec` to get a custom parameter name.
-        exec(
-            textwrap.dedent(
-                f"""
         @typechecker
-        def typecheck({field.name}: annotation):
+        def typecheck(x: annotation):
             pass
 
-        typecheck(value)
-        """
-            )
-        )
+        try:
+            typecheck(value)
+        except Exception as e:
+            if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
+                shape_info = _exc_shape_info(get_shape_memo())
+                if shape_info != "":
+                    msg = (
+                        "The above typechecking error occurred due to a mismatch "
+                        f"between the value and annotation for field '{field.name}' in "
+                        "dataclass "
+                        f"'{self.__class__.__module__}.{self.__class__.__qualname__}'. "
+                    )
+                    e.add_note(_jaxtyping_note_str(_spacer + msg + shape_info))
+            raise
 
 
 def _jaxtyped_typechecker(typechecker):
@@ -239,3 +258,45 @@ def _jaxtyped_typechecker(typechecker):
         return kls
 
     return _wrapper
+
+
+def _no_jaxtyping_note(e):
+    try:
+        notes = e.__notes__
+    except AttributeError:
+        return True
+    else:
+        for note in notes:
+            if isinstance(note, _jaxtyping_note_str):
+                return False
+        return True
+
+
+class _jaxtyping_note_str(str):
+    pass
+
+
+_spacer = "--------------------\n"
+
+
+def _exc_shape_info(memos) -> str:
+    single_memo, variadic_memo, pytree_memo = memos
+    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}")
+    if len(pytree_memo) > 0:
+        pieces.append(
+            "The current values for each jaxtyping pytree structure annotation are as "
+            "follows."
+        )
+        for name, structure in pytree_memo.items():
+            pieces.append(f"{name}={structure}")
+    return "\n".join(pieces)

jaxtyping/_pytree_type.py

@@ -80,7 +80,7 @@ class _MetaPyTree(type):
 
         def is_leaftype(x, new_scope=True):
             if new_scope and cls.structure is not None:
-                set_treepath_memo(-1, "")
+                set_treepath_memo(None, cls.structure)
             try:
                 accepts_leaftype(x)
             except _TypeCheckError:

jaxtyping/_storage.py

@@ -18,6 +18,7 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import threading
+from typing import Optional
 
 
 _shape_storage = threading.local()
@@ -45,13 +46,15 @@ def set_shape_memo(single_memo, variadic_memo, pytree_memo) -> None:
         _shape_storage.memo_stack[-1] = single_memo, variadic_memo, pytree_memo
 
 
-def push_shape_memo() -> None:
+def push_shape_memo():
     try:
         memo_stack = _shape_storage.memo_stack
     except AttributeError:
         # Can't be done when `_stack_storage` is created for reasons I forget.
         memo_stack = _shape_storage.memo_stack = []
-    memo_stack.append(({}, {}, {}))
+    memos = ({}, {}, {})
+    memo_stack.append(memos)
+    return memos
 
 
 def pop_shape_memo() -> None:
@@ -65,14 +68,18 @@ def clear_treepath_memo() -> None:
     _treepath_storage.value = None
 
 
-def set_treepath_memo(index: int, structure: str) -> None:
+def set_treepath_memo(index: Optional[int], structure: str) -> None:
     if hasattr(_treepath_storage, "value") and _treepath_storage.value is not None:
         raise ValueError(
             "Cannot typecheck annotations of the form "
             "`PyTree[PyTree[Shaped[Array, '?foo'], 'T'], 'S']` as it is ambiguous "
             "which PyTree the `?` annotation refers to."
         )
-    _treepath_storage.value = str(index) + structure
+    if index is None:
+        _treepath_storage.value = f"~~delete~~({structure}) "
+    else:
+        # Appears in error messages, so human-readable
+        _treepath_storage.value = f"(Leaf {index} in structure {structure}) "
 
 
 def get_treepath_memo() -> str: