diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index 1d907f6..f857c1a 100644 --- a/jaxtyping/_decorator.py +++ b/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) diff --git a/jaxtyping/_pytree_type.py b/jaxtyping/_pytree_type.py index e0a4aec..10ebcdf 100644 --- a/jaxtyping/_pytree_type.py +++ b/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: diff --git a/jaxtyping/_storage.py b/jaxtyping/_storage.py index 63cac07..648d70a 100644 --- a/jaxtyping/_storage.py +++ b/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: