diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 177a970..991f1b1 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -30,7 +30,11 @@ from ._array_types import ( set_array_name_format as set_array_name_format, ) from ._config import config as config -from ._decorator import jaxtyped as jaxtyped, TypeCheckError as TypeCheckError +from ._decorator import jaxtyped as jaxtyped +from ._errors import ( + AnnotationError as AnnotationError, + TypeCheckError as TypeCheckError, +) from ._import_hook import install_import_hook as install_import_hook from ._ipython_extension import load_ipython_extension as load_ipython_extension diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py index 5ac236f..417dc69 100644 --- a/jaxtyping/_array_types.py +++ b/jaxtyping/_array_types.py @@ -27,7 +27,7 @@ from typing import Any, Literal, NoReturn, Optional, Union import numpy as np -from ._raise import jaxtyping_raise, jaxtyping_raise_from +from ._errors import AnnotationError from ._storage import ( get_shape_memo, get_treeflatten_memo, @@ -133,16 +133,13 @@ def _check_dims( # Make a copy to avoid `__builtins__` getting added as a key. eval_size = eval(elem, single_memo.copy()) except NameError as e: - jaxtyping_raise_from( - NameError( - 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, - ) + raise AnnotationError( + 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." + ) from e if eval_size != obj_size: return False else: @@ -180,8 +177,8 @@ class _MetaAbstractArray(type): if len(repr_dtype) == 2 and repr_dtype[0] == "torch": dtype = repr_dtype[1] else: - jaxtyping_raise( - RuntimeError("Unrecognised array/tensor type to extract dtype from") + raise AnnotationError( + "Unrecognised array/tensor type to extract dtype from" ) if cls.dtypes is not _any_dtype: @@ -561,12 +558,10 @@ def _make_array(array_type, dim_str, dtypes, name): class _MetaAbstractDtype(type): def __instancecheck__(cls, obj: Any) -> NoReturn: - jaxtyping_raise( - RuntimeError( - f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to " - "check just the dtype of an array, then use " - f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.' - ) + raise AnnotationError( + f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to " + "check just the dtype of an array, then use " + f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.' ) def __getitem__(cls, item: tuple[Any, str]): diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index a36fb88..401aee5 100644 --- a/jaxtyping/_decorator.py +++ b/jaxtyping/_decorator.py @@ -35,16 +35,10 @@ else: from ._config import config +from ._errors import AnnotationError, TypeCheckError from ._storage import pop_shape_memo, push_shape_memo -class TypeCheckError(TypeError): - pass - - -TypeCheckError.__module__ = "jaxtyping" # appears in error messages - - class _Sentinel: def __repr__(self): return "sentinel" @@ -381,30 +375,29 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel): # called. try: param_fn(*args, **kwargs) + except AnnotationError: + raise except Exception as e: - if hasattr(e, "_jaxtyping_malformed"): - raise + argmsg = _get_problem_arg( + param_signature, args, kwargs, module, typechecker + ) + try: + name = fn.__name__ + except AttributeError: + name = fn.__class__.__name__ + 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 arguments: {param_values}\n" + f"Parameter annotations: {param_hints}.\n" + + _exc_shape_info(memos) + ) + if config.jaxtyping_remove_typechecker_stack: + raise TypeCheckError(msg) from None else: - argmsg = _get_problem_arg( - param_signature, args, kwargs, module, typechecker - ) - try: - name = fn.__name__ - except AttributeError: - name = fn.__class__.__name__ - 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 arguments: {param_values}\n" - f"Parameter annotations: {param_hints}.\n" - + _exc_shape_info(memos) - ) - if config.jaxtyping_remove_typechecker_stack: - raise TypeCheckError(msg) from None - else: - raise TypeCheckError(msg) from e + raise TypeCheckError(msg) from e # Actually call the function. out = fn(*args, **kwargs) @@ -427,39 +420,36 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel): kwargs[output_name] = out try: full_fn(*args, **kwargs) + except AnnotationError: + raise except Exception as e: - if hasattr(e, "_jaxtyping_malformed"): - raise + try: + name = fn.__name__ + except AttributeError: + name = fn.__class__.__name__ + 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 return_hint.startswith( + ""): + return_hint = return_hint[8:-2] + msg = ( + "Type-check error whilst checking the return value " + f"of {name}.\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) + ) + if config.jaxtyping_remove_typechecker_stack: + raise TypeCheckError(msg) from None else: - try: - name = fn.__name__ - except AttributeError: - name = fn.__class__.__name__ - 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 return_hint.startswith( - ""): - return_hint = return_hint[8:-2] - msg = ( - "Type-check error whilst checking the return value " - f"of {name}.\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) - ) - if config.jaxtyping_remove_typechecker_stack: - raise TypeCheckError(msg) from None - else: - raise TypeCheckError(msg) from e + raise TypeCheckError(msg) from e return out finally: diff --git a/jaxtyping/_errors.py b/jaxtyping/_errors.py new file mode 100644 index 0000000..918917f --- /dev/null +++ b/jaxtyping/_errors.py @@ -0,0 +1,12 @@ +class TypeCheckError(TypeError): + pass + + +# Not inheriting from TypeError as that gets caught and re-reraised as just a TypeError +# when using typeguard<3. +class AnnotationError(Exception): + pass + + +TypeCheckError.__module__ = "jaxtyping" +AnnotationError.__module__ = "jaxtyping" diff --git a/jaxtyping/_pytree_type.py b/jaxtyping/_pytree_type.py index df89997..b95bee8 100644 --- a/jaxtyping/_pytree_type.py +++ b/jaxtyping/_pytree_type.py @@ -24,7 +24,7 @@ from typing import Any, Generic, TypeVar import jax.tree_util as jtu import typeguard -from ._raise import jaxtyping_raise_from +from ._errors import AnnotationError from ._storage import ( clear_treeflatten_memo, clear_treepath_memo, @@ -141,14 +141,11 @@ class _MetaPyTree(type): try: prev_structure = pytree_memo[identifier] except KeyError as e: - jaxtyping_raise_from( - NameError( - f"Cannot process composite structure '{cls.structure}' " - f"as the structure name {identifier} has not been seen " - "before." - ), - e, - ) + raise AnnotationError( + f"Cannot process composite structure '{cls.structure}' " + f"as the structure name {identifier} has not been seen " + "before." + ) from e # Not using `PyTreeDef.compose` due to JAX bug #18218. prev_pytree = jtu.tree_unflatten( prev_structure, [0] * prev_structure.num_leaves diff --git a/jaxtyping/_raise.py b/jaxtyping/_raise.py deleted file mode 100644 index 38d3f1a..0000000 --- a/jaxtyping/_raise.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import NoReturn - - -def jaxtyping_raise(e) -> NoReturn: - """Raises `e`, whilst adding a tag that it should not be intercepted by - `TypeCheckError`. All `raise` statements from within `__instancecheck__` should use - this. - """ - __tracebackhide__ = True - try: - raise e - except Exception as f: - f._jaxtyping_malformed = True - raise - - -def jaxtyping_raise_from(e, e_base) -> NoReturn: - __tracebackhide__ = True - try: - raise e from e_base - except Exception as f: - f._jaxtyping_malformed = True - raise diff --git a/jaxtyping/_storage.py b/jaxtyping/_storage.py index 73d1f64..7cd1a30 100644 --- a/jaxtyping/_storage.py +++ b/jaxtyping/_storage.py @@ -20,7 +20,7 @@ import threading from typing import Any, Optional -from ._raise import jaxtyping_raise +from ._errors import AnnotationError _shape_storage = threading.local() @@ -80,12 +80,10 @@ def clear_treepath_memo() -> None: def set_treepath_memo(index: Optional[int], structure: str) -> None: if hasattr(_treepath_storage, "value") and _treepath_storage.value is not None: - jaxtyping_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." - ) + raise AnnotationError( + "Cannot typecheck annotations of the form " + "`PyTree[PyTree[Shaped[Array, '?foo'], 'T'], 'S']` as it is ambiguous " + "which PyTree the `?` annotation refers to." ) if index is None: _treepath_storage.value = f"~~delete~~({structure}) " @@ -96,12 +94,10 @@ def set_treepath_memo(index: Optional[int], structure: str) -> None: def get_treepath_memo() -> str: if not hasattr(_treepath_storage, "value") or _treepath_storage.value is None: - jaxtyping_raise( - ValueError( - "Cannot use `?` annotations, e.g. `Shaped[Array, '?foo']`, except " - "when contained with structured `PyTree` annotations, e.g. " - "`PyTree[Shaped[Array, '?foo'], 'T']`." - ) + raise AnnotationError( + "Cannot use `?` annotations, e.g. `Shaped[Array, '?foo']`, except " + "when contained with structured `PyTree` annotations, e.g. " + "`PyTree[Shaped[Array, '?foo'], 'T']`." ) return _treepath_storage.value diff --git a/test/test_array.py b/test/test_array.py index 1b81cae..de660cf 100644 --- a/test/test_array.py +++ b/test/test_array.py @@ -29,6 +29,7 @@ import torch from jaxtyping import ( AbstractDtype, + AnnotationError, Array, ArrayLike, Bool, @@ -448,7 +449,7 @@ def test_incomplete_symbolic(jaxtyp, typecheck, getkey): pass x = jr.normal(getkey(), (4,)) - with pytest.raises(NameError): + with pytest.raises(AnnotationError): foo(x) diff --git a/test/test_pytree.py b/test/test_pytree.py index a1b416e..9c714bd 100644 --- a/test/test_pytree.py +++ b/test/test_pytree.py @@ -26,7 +26,7 @@ import jax.random as jr import pytest import jaxtyping -from jaxtyping import Array, Float, PyTree +from jaxtyping import AnnotationError, Array, Float, PyTree from .helpers import make_mlp, ParamError @@ -330,7 +330,7 @@ def test_treepath_dependence_missing_structure_annotation(jaxtyp, typecheck, get x1 = jr.normal(getkey(), (2,)) y1 = jr.normal(getkey(), (2,)) - with pytest.raises(ValueError, match="except when contained with structured"): + with pytest.raises(AnnotationError, match="except when contained with structured"): f(x1, y1) @@ -340,5 +340,5 @@ def test_treepath_dependence_multiple_structure_annotation(jaxtyp, typecheck, ge pass x1 = jr.normal(getkey(), (2,)) - with pytest.raises(ValueError, match="ambiguous which PyTree"): + with pytest.raises(AnnotationError, match="ambiguous which PyTree"): f(x1)