wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-07 17:00:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

Simplified internals by removing jaxtyping_raise; jaxtyping_malformed.

Browse Source
This commit is contained in:
Patrick Kidger
2023-12-05 19:06:00 -08:00
parent e5cc75e4a3
commit 33cf4fcdac
9 changed files with 101 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jaxtyping/__init__.py
+5 -1
View File
@@ -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
jaxtyping/_array_types.py
+14 -19
View File
@@ -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]):
jaxtyping/_decorator.py
+50 -60
View File
@@ -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(
"<class '"
) 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 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(
"<class '"
) 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 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:
jaxtyping/_errors.py
+12
View File
@@ -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"
jaxtyping/_pytree_type.py
+6 -9
View File
@@ -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
jaxtyping/_raise.py
-23
View File
@@ -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
jaxtyping/_storage.py
+9 -13
View File
@@ -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
test/test_array.py
+2 -1
View File
@@ -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)
test/test_pytree.py
+3 -3
View File
@@ -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)