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

Pretty error messages: fixes #6.

Browse Source 
Phew, this ended up being a pretty complicated change!
The basic summary is that we now support the syntax
```
@jaxtyped(typechecker=typechecker)
def f(...): ...
```
and when using this, we now get pretty error messages about what went
wrong.

(
The old syntax, i.e.
```
@jaxtyped
@typechecker
def f(...): ...
```
is still supported, but doesn't give much information.
)

The internals of this do quite a lot of magic! In particular we
dynamically create quite a lot of functions and test the provided
arguments against their signatures. The overhead should still be
minimal under `jax.jit`, though.
(TODO: what's the overhead like in non-jit situations, e.g. PyTorch?
I've tried to minimise the overhead throughout just to be sure, but
perhaps PyTorch users should stick to the old syntax?)
This commit is contained in:
Patrick Kidger
2023-11-27 09:50:02 -08:00
parent 63e0fdff74
commit 12d540794f
14 changed files with 984 additions and 389 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/api/array.md
+1 -1
View File
@@ -24,7 +24,7 @@ In addition some modifiers can be applied:
`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 _ _"`.
- 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.md).)
- Prepend `?` to a dimension to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree/).)
When using multiple modifiers, their order does not matter.
docs/api/runtime-type-checking.md
+26 -9
View File
@@ -2,27 +2,40 @@
(See the [FAQ](../faq.md) for details on static type checking.)
Runtime type checking **synergises beautifully with `jax.jit`!** All shape checks will be performed at trace-time only, and will not impact runtime performance.
Runtime type checking **synergises beautifully with `jax.jit`!** All shape checks will be performed only whilst tracing, and will not impact runtime performance.
Runtime type-checking should be performed using a library like [typeguard](https://github.com/agronholm/typeguard) or [beartype](https://github.com/beartype/beartype).
There are two approaches: either use [`jaxtyping.jaxtyped`][] to typecheck a single function, or [`jaxtyping.install_import_hook`][] to typecheck a whole codebase.
The types provided by `jaxtyping`, e.g. `Float[Array, "batch channels"]`, are all compatible with `isinstance` checks, e.g. `isinstance(x, Float[Array, "batch channels"])`. This means that jaxtyping should be compatible with all runtime type checkers out-of-the-box.
In either case, the actual business of checking types is performed with the help of a runtime type-checking library. The two most popular are [beartype](https://github.com/beartype/beartype) and [typeguard](https://github.com/agronholm/typeguard). (If using typeguard, then specifically the version `2.*` series should be used. Later versions -- `3` and `4` -- have some known issues.)
Some additional context is needed to ensure consistency between multiple argments (i.e. that shapes match up between arrays). For this, you can use either `jaxtyping.jaxtyped` to add this capability to a single function, or `jaxtyping.install_import_hook` to add this capability to a whole codebase. If either are too much magic for you, you can safely use neither and have just single-argument type checking.
---
::: jaxtyping.jaxtyped
---
It can be a lot of effort to add `@jaxtyped` decorators all over your codebase.
(Not to mention that double-decorators everywhere are a bit ugly.)
The easier option is usually to use the import hook.
::: jaxtyping.install_import_hook
---
#### Pytest hook
The import hook can be installed at test-time only, as a pytest hook. From the command line the syntax is:
```
pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
```
or in `pyproject.toml`:
```toml
[tool.pytest.ini_options]
addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
```
or in `pytest.ini`:
```ini
[pytest]
addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
```
This example will apply the import hook to all modules whose names start with either `foo` or `bar.baz`. The typechecker used in this example is `beartype.beartype`.
#### IPython extension
If you are running in an IPython environment (for example a Jupyter or Colab notebook), then the jaxtyping hook can be automatically ran via a custom magic:
@@ -32,3 +45,7 @@ import jaxtyping
%jaxtyping.typechecker beartype.beartype # or any other runtime type checker
```
Place this at the start of your notebook -- everything that is directly defined in the notebook, after this magic is run, will be hook'd.
#### Other runtime type-checking libraries
Beartype and typeguard happen to be the two most popular runtime type-checking libraries (at least at time of writing), but jaxtyping should be compatible with all runtime type checkers out-of-the-box. The runtime type-checking library just needs to provide a type-checking decorator (analgous to `beartype.beartype` or `typeguard.typechecked`), and perform `isinstance` checks against jaxtyping's types.
jaxtyping/__init__.py
+1 -1
View File
@@ -29,7 +29,7 @@ from ._array_types import (
has_jax,
set_array_name_format as set_array_name_format,
)
from ._decorator import jaxtyped as jaxtyped
from ._decorator import jaxtyped as jaxtyped, 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
+20 -13
View File
@@ -27,6 +27,7 @@ from typing import Any, Literal, NoReturn, Optional, Union
import numpy as np
from ._raise import jaxtyping_raise, jaxtyping_raise_from
from ._storage import get_shape_memo, get_treepath_memo, set_shape_memo
@@ -124,13 +125,16 @@ def _check_dims(
# Make a copy to avoid `__builtins__` getting added as a key.
eval_size = eval(cls_dim.expr, single_memo.copy())
except NameError as e:
raise 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 "
"arguments."
) from 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 "
"arguments."
),
e,
)
if eval_size != obj_size:
return False
else:
@@ -186,8 +190,8 @@ class _MetaAbstractArray(type):
if len(repr_dtype) == 2 and repr_dtype[0] == "torch":
dtype = repr_dtype[1]
else:
raise RuntimeError(
"Unrecognised array/tensor type to extract dtype from"
jaxtyping_raise(
RuntimeError("Unrecognised array/tensor type to extract dtype from")
)
if cls.dtypes is not _any_dtype:
@@ -553,10 +557,12 @@ def _make_array(array_type, dim_str, dtypes, name):
class _MetaAbstractDtype(type):
def __instancecheck__(cls, obj: Any) -> NoReturn:
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, "..."]`.'
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, "..."]`.'
)
)
def __getitem__(cls, item: tuple[Any, str]):
@@ -567,6 +573,7 @@ class _MetaAbstractDtype(type):
"Ellipsis can be used to accept any shape: `Float[Array, '...']`."
)
array_type, dim_str = item
dim_str = dim_str.strip()
del item
if typing.get_origin(array_type) in _union_types:
out = [
jaxtyping/_decorator.py
+532 -153
View File
@@ -20,10 +20,11 @@
import dataclasses
import functools as ft
import inspect
import itertools as it
import sys
import types
import weakref
from typing import get_args, get_origin
from typing import Any, get_args, get_origin, get_type_hints, overload
try:
@@ -34,113 +35,340 @@ else:
traceback_util.register_exclusion(__file__)
from ._storage import get_shape_memo, pop_shape_memo, push_shape_memo
from ._storage import pop_shape_memo, push_shape_memo
_jaxtyped_fns = weakref.WeakSet()
def jaxtyped(fn):
"""Used in conjunction with a runtime type checker. Decorate a function with this to
have shapes checked for consistency across multiple arguments.
class TypeCheckError(TypeError):
pass
Note that `@jaxtyped` is applied above the type checker.
TypeCheckError.__module__ = "jaxtyping" # appears in error messages
@overload
def jaxtyped(*, typechecker=None):
...
@overload
def jaxtyped(fn, *, typechecker=None):
...
def jaxtyped(fn=None, *, typechecker=None):
"""Decorate a function with this to perform runtime type-checking of its arguments
and return value. Decorate a dataclass to perform type-checking of its attributes.
!!! Example
```python
# Import both the annotation and the `jaxtyped` decorator from `jaxtyping`
from jaxtyping import Array, Float32, jaxtyped
from jaxtyping import Array, Float, jaxtyped
# Use your favourite typechecker: usually one of the two lines below.
from typeguard import typechecked as typechecker
from beartype import beartype as typechecker
# Write your function. @jaxtyped must be applied above @typechecker!
@jaxtyped
@typechecker
def batch_outer_product(x: Float32[Array, "b c1"],
y: Float32[Array, "b c2"]
) -> Float32[Array, "b c1 c2"]:
# Type-check a function
@jaxtyped(typechecker=typechecker)
def batch_outer_product(x: Float[Array, "b c1"],
y: Float[Array, "b c2"]
) -> Float[Array, "b c1 c2"]:
return x[:, :, None] * y[:, None, :]
# Type-check a dataclass
@jaxtyped(typechecker=typechecker)
@dataclass
class MyDataclass:
x: int
y: Float[Array "b c"]
```
**Arguments:**
- `fn`: The function or dataclass to decorate.
- `typechecker`: The runtime type-checker to use. This should be a function
decorator that will raise an exception if there is a type error, e.g.
```python
@typechecker
def f(x: int):
pass
f("a string is not an integer") # this line should raise an exception
```
Common choices are `typechecker=beartype.beartype` or
`typechecker=typeguard.typechecked`. Can also be set as `typechecker=None` to
skip automatic runtime type-checking, but still support manual `isinstance`
checks inside the function body:
```python
@jaxtyped
def f(x):
assert isinstance(x, Float[Array, "batch channel"])
```
**Returns:**
If `fn` is a function (including a `staticmethod`, `classmethod`, or `property`),
then a wrapped function is returned.
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, all `isinstance` shape checks are scoped to the thread-local dynamic
context of a `jaxtyped` call. A new dynamic context will allow different dimensions
sizes to be bound to the same name. After this new dynamic context is finished
then the old one is returned to.
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.
For example, this means you could leave off the `@jaxtyped` decorator to enforce
that this function use the same axis sizes as the function it was called from.
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`.
Likewise, 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.)
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 axis name-size
pairs; those that fail will not. As such it is safe to write e.g.
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
syntax:
```python
@jaxtyped
@typechecker
def f(...): ...
```
This is still supported, but the `jaxtyped(typechecker=typechecker)` syntax
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.
"""
if type(fn) is types.FunctionType and fn in _jaxtyped_fns:
if fn is None:
return ft.partial(jaxtyped, typechecker=typechecker)
elif type(fn) is types.FunctionType and fn in _jaxtyped_fns:
return fn
elif inspect.isclass(fn):
if dataclasses.is_dataclass(fn) and typechecker is not None:
# This does not check that the arguments passed to `__init__` match the
# type annotations. There may be a custom user `__init__`, or a
# dataclass-generated `__init__` used alongside
# `equinox.field(converter=...)`
init = fn.__init__
@ft.wraps(init)
def __init__(self, *args, **kwargs):
init(self, *args, **kwargs)
# `fn.__init__` is late-binding to the `__init__` function that
# we're in now. (Or to someone else's monkey-patch.) Either way,
# this checks that we're in the "top-level" `__init__`, and not one
# that is being called via `super()`. We don't want to trigger too
# early, before all fields have been assigned.
#
# We're not checking `if self.__class__ is fn` because Equinox
# replaces the with a defrozen version of itself during `__init__`,
# so the check wouldn't trigger.
#
# We're not doing this check by adding it to the end of the
# metaclass `__call__`, because Python doesn't allow you
# monkey-patch metaclasses.
if self.__class__.__init__ is fn.__init__:
_check_dataclass_annotations(self, typechecker)
fn.__init__ = __init__
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
else:
raise ValueError(
"jaxtyped may only be added as a class decorator to dataclasses"
)
# It'd be lovely if we could handle arbitrary descriptors, and not just the builtin
# ones. Unfortunately that means returning a class instance with a __get__ method,
# and that turns out to break loads of other things. See beartype issue #211 and
# jaxtyping issue #71.
elif isinstance(fn, classmethod):
return classmethod(jaxtyped(fn.__func__))
return classmethod(jaxtyped(fn.__func__, typechecker=typechecker))
elif isinstance(fn, staticmethod):
return staticmethod(jaxtyped(fn.__func__))
return staticmethod(jaxtyped(fn.__func__, typechecker=typechecker))
elif isinstance(fn, property):
if fn.fget is None:
fget = None
else:
fget = jaxtyped(fn.fget)
fget = jaxtyped(fn.fget, typechecker=typechecker)
if fn.fset is None:
fset = None
else:
fset = jaxtyped(fn.fset)
fset = jaxtyped(fn.fset, typechecker=typechecker)
if fn.fdel is None:
fdel = None
else:
fdel = jaxtyped(fn.fdel)
fdel = jaxtyped(fn.fdel, typechecker=typechecker)
return property(fget=fget, fset=fset, fdel=fdel)
elif fn == "context":
return _JaxtypingContext()
else:
if typechecker is None:
# Probably being used in the old style as
# ```
# @jaxtyped
# @typechecker
# def foo(x: int): ...
# ```
# in which case make a best-effort attempt to add shape information for any
# type errors.
@ft.wraps(fn)
def wrapped_fn(*args, **kwargs): # pyright: ignore
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()
@ft.wraps(fn)
def wrapped_fn(*args, **kwargs):
memos = push_shape_memo()
else:
# New-style
# ```
# @jaxtyped(typechecker=typechecker)
# def foo(x: int): ...
# ```
# in which case we can do a better job reporting errors.
full_signature = inspect.signature(fn)
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()
destring_annotations = get_type_hints(fn, include_extras=True)
except NameError:
# Best-effort attempt to destringify annotations.
pass
else:
new_params = []
for p_name, p_value in full_signature.parameters.items():
p_annotation = destring_annotations.get(p_name, p_value.annotation)
p_value = p_value.replace(annotation=p_annotation)
new_params.append(p_value)
return_annotation = destring_annotations.get(
"return", full_signature.return_annotation
)
full_signature = full_signature.replace(
parameters=new_params, return_annotation=return_annotation
)
param_signature = full_signature.replace(
return_annotation=inspect.Signature.empty
)
module = getattr(fn, "__module__", "generated")
full_fn, output_name = _make_fn_with_signature(
"check_return", full_signature, module, output=True
)
full_fn = typechecker(full_fn)
param_fn = _make_fn_with_signature(
"check_params", param_signature, module, output=False
)
param_fn = typechecker(param_fn)
@ft.wraps(fn)
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)
memos = push_shape_memo()
try:
# First type-check just the parameters before the function is
# called.
try:
param_fn(*args, **kwargs)
except Exception as e:
if hasattr(e, "_jaxtyping_malformed"):
raise
else:
argmsg = _get_problem_arg(
param_signature, args, kwargs, module, typechecker
)
try:
name = fn.__name__
except AttributeError:
name = fn.__class__.__name__
paramstr = _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"
+ _exc_shape_info(memos)
)
raise TypeCheckError(msg) from e
# Actually call the function.
out = fn(*args, **kwargs)
if full_signature.return_annotation is not inspect.Signature.empty:
# Now type-check the return value. We need to include the
# parameters in the type-checking here in case there are any
# type variables shared across the parameters and return.
#
# Incidentally this does mean that if `fn` mutates its arguments
# so that they no longer satisfy their type annotations, this
# will throw an error here. But that's like, super weird, so
# don't do that. An error in that scenario is probably still
# desirable.
#
# There is a small performance concern here when used in
# non-jit'd contexts, like PyTorch, due to the duplicate
# checking of the parameters. Unfortunately there doesn't seem
# to be a way around that, so c'est la vie.
kwargs[output_name] = out
try:
full_fn(*args, **kwargs)
except Exception as e:
if hasattr(e, "_jaxtyping_malformed"):
raise
else:
try:
name = fn.__name__
except AttributeError:
name = fn.__class__.__name__
paramstr = _remove_typing(param_signature)
returnstr = _remove_typing(
full_signature.return_annotation
)
if returnstr.startswith(
"<class '"
) and returnstr.endswith("'>"):
returnstr = returnstr[8:-2]
kwargs.pop(output_name)
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"
+ _exc_shape_info(memos)
)
raise TypeCheckError(msg) from e
return out
finally:
pop_shape_memo()
_jaxtyped_fns.add(wrapped_fn)
return wrapped_fn
@@ -154,24 +382,19 @@ class _JaxtypingContext:
pop_shape_memo()
@jaxtyped
def _check_dataclass_annotations(self, typechecker):
"""Creates and calls a function that checks the attributes of `self`
`self` should be a dataclass instancae. `typechecker` should be e.g.
`beartype.beartype` or `typeguard.typechecked`.
"""
parameters = []
values = {}
for field in dataclasses.fields(self):
for kls in self.__class__.__mro__:
try:
annotation = kls.__annotations__[field.name]
except KeyError:
pass
else:
break
else:
raise TypeError
annotation = field.type
if isinstance(annotation, str):
# Don't support stringified annotations. These are basically impossible to
# Don't check stringified annotations. These are basically impossible to
# resolve correctly, so just skip them.
# This does mean that annotations like `type["Foo"]` will just fail. There
# doesn't seem to be any way to even detect a partially-stringified
# annotation.
continue
if get_origin(annotation) is type:
args = get_args(annotation)
@@ -186,100 +409,236 @@ def _check_dataclass_annotations(self, typechecker):
except AttributeError:
continue # allow uninitialised fields, which are allowed on dataclasses
@typechecker
def typecheck(x: annotation):
pass
parameters.append(
inspect.Parameter(
field.name,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=field.type,
)
)
values[field.name] = 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
signature = inspect.Signature(parameters)
module = self.__class__.__module__
f = _make_fn_with_signature(
self.__class__.__name__, signature, module, output=False
)
f = jaxtyped(f, typechecker=typechecker)
f(**values)
def _jaxtyped_typechecker(typechecker):
"""A decorator added by the import hook to all classes. Only affects dataclasses.
def _make_fn_with_signature(
name: str, signature: inspect.Signature, module: str, output: bool
):
"""Dynamically creates a function `fn` with name `name` and signature `signature`.
Will be called as
```
@_jaxtyped_typechecker(beartype.beartype)
@dataclasses.dataclass
class SomeDataclass:
...
```
If `output=True` then `fn` will consume an additional keyword-only argument (in
addition to the provided signature), and will directly return this argument. In this
case the returned value from `_make_fn_with_signature` is a 2-tuple of `(fn, name)`,
where `fn` is the generated function, and `name` is the name of this extra argument.
After initialisation, this will check that all fields of the dataclass match their
specified type annotation.
If `output=False` then `fn` will just have a single `pass` statement, and the
returned value from `_make_fn_with_signature` will just be `fn`.
---
Note that this function operates by dynamically creating and eval'ing a string, not
simply by assigning `__signature__` and `__annotations__`. The latter is enough for
typeguard (at least v2), but does not work with beartype (at least v16).
"""
# typechecker is expected to probably be either `typeguard.typechecked`, or
# `beartype.beartype`, or `None`.
pos = []
pos_or_key = []
varpos = []
key = []
varkey = []
for p in signature.parameters.values():
if p.kind == inspect.Parameter.POSITIONAL_ONLY:
pos.append(p)
elif p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
pos_or_key.append(p)
elif p.kind == inspect.Parameter.VAR_POSITIONAL:
varpos.append(p)
elif p.kind == inspect.Parameter.KEYWORD_ONLY:
key.append(p)
elif p.kind == inspect.Parameter.VAR_KEYWORD:
varkey.append(p)
else:
assert False
if typechecker is None:
typechecker = lambda x: x
def _wrapper(kls):
assert inspect.isclass(kls)
if dataclasses.is_dataclass(kls):
# This does not check that the arguments passed to `__init__` match the
# type annotations. There may be a custom user `__init__`, or a
# dataclass-generated `__init__` used alongside
# `equinox.field(converter=...)`
init = kls.__init__
@ft.wraps(init)
def __init__(self, *args, **kwargs):
init(self, *args, **kwargs)
# `kls.__init__` is late-binding to the `__init__` function that we're
# in now. (Or to someone else's monkey-patch.) Either way, this checks
# that we're in the "top-level" `__init__`, and not one that is being
# called via `super()`. We don't want to trigger too early, before all
# fields have been assigned.
#
# We're not checking `if self.__class__ is kls` because Equinox replaces
# the with a defrozen version of itself during `__init__`, so the check
# wouldn't trigger.
#
# We're not doing this check by adding it to the end of the metaclass
# `__call__`, because Python doesn't allow you monkey-patch metaclasses.
if self.__class__.__init__ is kls.__init__:
_check_dataclass_annotations(self, typechecker)
kls.__init__ = __init__
return kls
return _wrapper
def _no_jaxtyping_note(e):
try:
notes = e.__notes__
except AttributeError:
return True
param_names = frozenset(signature.parameters.keys())
if output:
output_name = _gensym(param_names, prefix="ret")
outstr = "return " + output_name
param_names = param_names | frozenset({output_name})
key.append(inspect.Parameter(output_name, kind=inspect.Parameter.KEYWORD_ONLY))
else:
for note in notes:
if isinstance(note, _jaxtyping_note_str):
return False
return True
outstr = "pass"
scope = {name: None}
name_to_annotation = {}
name_to_default = {}
param_triples = (
(p.name, p.annotation, p.default) for p in signature.parameters.values()
)
if output:
triples = it.chain(
param_triples,
[
("return", signature.return_annotation, inspect.Signature.empty),
(output_name, Any, inspect.Signature.empty),
],
)
else:
triples = it.chain(
param_triples,
[("return", signature.return_annotation, inspect.Signature.empty)],
)
for p_name, p_annotation, p_default in triples:
annotation_name = _gensym(frozenset(scope.keys()) | param_names, prefix="T")
name_to_annotation[p_name] = annotation_name
if p_annotation is inspect.Signature.empty or isinstance(p_annotation, str):
# If we have a stringified annotation here it's because the get_type_hints
# lookup above failed. Typically this occurs when using a local variable as
# the annotation. In this case we really have no idea what the annotation
# refers to, so just set it to Any.
# This does mean that we don't handle partially-stringified local
# annotations, e.g. `type["Foo"]` for some local type `Foo`. Those will
# probably just error out. Nothing better we can do about that
# unfortunately.
scope[annotation_name] = Any
else:
scope[annotation_name] = p_annotation
default_name = _gensym(frozenset(scope.keys()) | param_names, prefix="default")
name_to_default[p_name] = default_name
scope[default_name] = p_default
argstr_pieces = []
if len(pos) > 0:
for p in pos:
argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
argstr_pieces.append("/")
if len(pos_or_key) > 0:
for p in pos_or_key:
argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
if len(varpos) == 1:
[p] = varpos
argstr_pieces.append(
"*" + _make_argpiece(p, name_to_annotation, name_to_default)
)
else:
assert len(varpos) == 0
if len(key) > 0:
argstr_pieces.append("*")
if len(key) > 0:
for p in key:
argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
if len(varkey) == 1:
[p] = varkey
argstr_pieces.append(
"**" + _make_argpiece(p, name_to_annotation, name_to_default)
)
else:
assert len(varkey) == 0
argstr = ", ".join(argstr_pieces)
if signature.return_annotation is inspect.Signature.empty:
retstr = ""
else:
retstr = f"-> {name_to_annotation['return']}"
fnstr = f"def {name}({argstr}){retstr}:\n {outstr}"
exec(fnstr, scope)
fn = scope[name]
fn.__module__ = module
assert fn is not None
if output:
return fn, output_name
else:
return fn
class _jaxtyping_note_str(str):
pass
def _gensym(names: frozenset[str], prefix: str) -> str:
assert prefix.isidentifier()
output_index = 0
output_name = prefix + str(output_index)
while output_name in names:
output_index += 1
output_name = prefix + str(output_index)
assert output_name.isidentifier()
return output_name
_spacer = "--------------------\n"
def _make_argpiece(p, name_to_annotation, name_to_default):
if p.default is inspect.Signature.empty:
return f"{p.name}: {name_to_annotation[p.name]}"
else:
return f"{p.name}: {name_to_annotation[p.name]} = {name_to_default[p.name]}"
def _get_problem_arg(
param_signature: inspect.Signature, args, kwargs, module, typechecker
) -> str:
"""Determines which argument was likely to be the problematic one responsible for
raising a type-check error.
"""
# No performance concerns, as this is only used when we're about to raise an error
# anyway.
for keep_name in param_signature.parameters.keys():
new_parameters = []
for p_name, p in param_signature.parameters.items():
if p_name == keep_name:
new_parameters.append(
inspect.Parameter(p.name, p.kind, annotation=p.annotation)
)
else:
new_parameters.append(inspect.Parameter(p.name, p.kind))
new_signature = inspect.Signature(new_parameters)
fn = _make_fn_with_signature(
"check_single_arg", new_signature, module, output=False
)
fn = typechecker(fn) # but no `jaxtyped`; keep the same environment.
try:
fn(*args, **kwargs)
except Exception:
return f"\nThe problem arose whilst typechecking argument '{keep_name}'."
else:
# Could not localise the problem to a single argument -- probably due to
# e.g. a mismatched typevar, which each individual argument is okay with.
return ""
def _remove_typing(x):
x = str(x)
x = x.replace(" jaxtyping.", " ")
x = x.replace("[jaxtyping.", "[")
x = x.replace("'jaxtyping.", "'")
x = x.replace(" typing.", " ")
x = x.replace("[typing.", "[")
x = x.replace("'typing.", "'")
return x
def _pformat(x):
# No performance concerns from delayed imports -- this is only used when we're about
# to raise an error anyway.
try:
# TODO(kidger): this is pretty ugly. We have a circular dependency
# equinox->jaxtyping->equinox. We could consider moving all the pretty-printing
# code from equinox into jaxtyping maybe? Or into some shared dependency?
import equinox as eqx
pformat = eqx.tree_pformat
except Exception:
import pprint
pformat = ft.partial(pprint.pformat, indent=2, compact=True)
return pformat(x)
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
pieces = []
if len(single_memo) > 0 or len(variadic_memo) > 0:
@@ -294,9 +653,29 @@ def _exc_shape_info(memos) -> str:
pieces.append(f"{name}={shape}")
if len(pytree_memo) > 0:
pieces.append(
"The current values for each jaxtyping pytree structure annotation are as "
"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)
class _jaxtyping_note_str(str):
"""Used with `_no_jaxtyping_note` to flag that a note came from jaxtyping."""
def _no_jaxtyping_note(e: Exception) -> bool:
"""Checks if any of the exception's notes are from jaxtyping."""
try:
notes = e.__notes__
except AttributeError:
return True
else:
for note in notes:
if isinstance(note, _jaxtyping_note_str):
return False
return True
_spacer = "--------------------\n"
jaxtyping/_import_hook.py
+47 -68
View File
@@ -78,18 +78,13 @@ def _optimized_cache_from_source(typechecker_hash, /, path, debug_override=None)
# changing the typechecker will hit a different cache.
# Version 7: Using the same md5 hash of the `typechecker` argument
# for importlib and decorator lookup.
# Version 8: Now using new-style `jaxtyped(typechecker=...)` rather than old-style
# double-decorators.
return cache_from_source(
path, debug_override, optimization=f"jaxtyping7{typechecker_hash}"
path, debug_override, optimization=f"jaxtyping8{typechecker_hash}"
)
def _dot_lookup(*elements):
out = ast.Name(id=elements[0], ctx=ast.Load())
for element in elements[1:]:
out = ast.Attribute(out, element, ctx=ast.Load())
return out
class Typechecker:
lookup = {}
@@ -130,7 +125,7 @@ class Typechecker:
if self.ast is None:
self.ast = (
ast.parse(
f"@jaxtyping._import_hook.Typechecker.lookup['{self.hash}']\n"
f"@jaxtyping.jaxtyped(typechecker=jaxtyping._import_hook.Typechecker.lookup['{self.hash}'])\n"
"def _():\n ..."
)
.body[0]
@@ -162,10 +157,9 @@ class JaxtypingTransformer(ast.NodeVisitor):
return node
def visit_ClassDef(self, node: ast.ClassDef):
func = _dot_lookup("jaxtyping", "_decorator", "_jaxtyped_typechecker")
node.decorator_list.insert(
0, ast.Call(func, [self._typechecker.get_ast()], keywords=[])
)
# Place at the start of the decorator list, so that `@dataclass` decorators get
# called first.
node.decorator_list.insert(0, self._typechecker.get_ast())
self._parents.append(node)
self.generic_visit(node)
self._parents.pop()
@@ -175,8 +169,13 @@ class JaxtypingTransformer(ast.NodeVisitor):
has_annotated_args = any(arg for arg in node.args.args if arg.annotation)
has_annotated_return = bool(node.returns)
if has_annotated_args or has_annotated_return:
# Place at the end of the decorator list, as otherwise we wrap e.g.
# `jax.custom_{jvp,vjp}` and lose the ability to `defjvp` etc.
# Place at the end of the decorator list, because:
# - as otherwise we wrap e.g. `jax.custom_{jvp,vjp}` and lose the ability
# to `defjvp` etc.
# - decorators frequently remove annotations from functions, and we'd like
# to use those annotations.
# - typeguard in particular wants to be at the end of the decorator list, as
# it works by recompling the wrapped function.
#
# Note that the counter-argument here is that we'd like to place this
# at the start of the decorator list, in case a typechecking annotation
@@ -184,14 +183,6 @@ class JaxtypingTransformer(ast.NodeVisitor):
# case we're just going to have to need to ask the user to remove their
# typechecking annotation (and let this decorator do it instead).
# It's more important we be compatible with normal JAX code.
#
#
# FWIW, typeguard also wants to be at the end of the decorator list, as it
# works by recompiling the wrapped function.
node.decorator_list.append(_dot_lookup("jaxtyping", "jaxtyped"))
# Place typechecker at the end of the decorator list, as decorators
# frequently remove annotations from functions and we'd like to
# use those annotations.
node.decorator_list.append(self._typechecker.get_ast())
self._parents.append(node)
@@ -290,7 +281,8 @@ class ImportHookManager:
# Deliberately no default for `typechecker` so that folks must opt-in to not having
# a typechecker.
def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optional[str]):
"""Automatically apply `@jaxtyped`, and optionally a type checker, as decorators.
"""Automatically apply the `@jaxtyped(typechecker=typechecker)` decorator to every
function and dataclass over a whole codebase.
!!! Tip "Usage"
@@ -298,18 +290,19 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
from jaxtyping import install_import_hook
# Plus any one of the following:
# decorate @jaxtyped and @typeguard.typechecked
# decorate `@jaxtyped(typechecker=typeguard.typechecked)`
with install_import_hook("foo", "typeguard.typechecked"):
import foo # Any module imported inside this `with` block, whose
import foo.bar # name begins with the specified string, will
import foo.bar.qux # automatically have both `@jaxtyped` and the specified
# typechecker applied to all of their functions.
# typechecker applied to all of their functions and
# dataclasses.
# decorate @jaxtyped and @beartype.beartype
# decorate `@jaxtyped(typechecker=beartype.beartype)`
with install_import_hook("foo", "beartype.beartype"):
...
# decorate only @jaxtyped (if you want that for some reason)
# decorate only `@jaxtyped` (if you want that for some reason)
with install_import_hook("foo", None):
...
```
@@ -326,22 +319,9 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
install_import_hook(["foo", "bar.baz"], ...)
```
The import hook will automatically decorate all functions, and check the attributes
assigned to dataclasses.
If the function already has any decorators on it, then both the `@jaxtyped` and the
typechecker decorators will get added at the bottom of the decorator list, e.g.
```python
@some_other_decorator
@jaxtyped
@beartype.beartype
def foo(...): ...
```
**Arguments:**:
- `modules`: the names of the modules in which to automatically apply `@jaxtyped`
and `@typechecked`.
- `modules`: the names of the modules in which to automatically apply `@jaxtyped`.
- `typechecker`: the module and function of the typechecker you want to use, as a
string. For example `typechecker="typeguard.typechecked"`, or
`typechecker="beartype.beartype"`. You may pass `typechecker=None` if you do not
@@ -351,7 +331,7 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
A context manager that uninstalls the hook on exit, or when you call `.uninstall()`.
??? Example "Example: end-user script"
!!! Example "Example: end-user script"
```python
### entry_point.py
@@ -366,7 +346,7 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
...
```
??? Example "Example: writing a library"
!!! Example "Example: writing a library"
```python
### __init__.py
@@ -378,30 +358,6 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
# so will be hook'd.
```
??? info "Pytest hook"
The import hook can be installed at test-time only, as a pytest hook. From the
command line the syntax is:
```
pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
```
or in `pyproject.toml`:
```toml
[tool.pytest.ini_options]
addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
```
or in `pytest.ini`:
```ini
[pytest]
addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
```
This example will apply the import hook to all modules whose names start with
either `foo` or `bar.baz`. The typechecker used in this example is
`beartype.beartype`.
(This is the author's preferred approach to performing runtime type-checking
with jaxtyping!)
!!! warning
Stringified dataclass annotations, e.g.
@@ -423,6 +379,29 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
x: tuple["int"]
```
will likely raise an error, and must not be used at all.
!!! warning
If a function already has any decorators on it, then `@jaxtyped` will get added
at the bottom of the decorator list, e.g.
```python
@some_other_decorator
@jaxtyped(typechecker=beartype.beartype)
def foo(...): ...
```
This is to support the common case in which
`some_other_decorator = jax.custom_jvp` etc.
If a class already has any decorators in it, then `@jaxtyped` will get added to
the top of the decorator list, e.g.
```python
@jaxtyped(typechecker=beartype.beartype)
@some_other_decorator
class A:
...
```
This is to support the common case in which
`some_other_decorator = dataclasses.dataclass`.
""" # noqa: E501
if isinstance(modules, str):
jaxtyping/_pytree_type.py
+43 -27
View File
@@ -19,11 +19,12 @@
import functools as ft
import typing
from typing import Generic, TypeVar
from typing import Any, Generic, TypeVar
import jax.tree_util as jtu
import typeguard
from ._raise import jaxtyping_raise_from
from ._storage import (
clear_treepath_memo,
get_shape_memo,
@@ -68,30 +69,41 @@ class _MetaPyTree(type):
return False
def _check(cls, obj, pytree_memo):
# We could use `isinstance` here but that would fail for more complicated
# types, e.g. PyTree[tuple[int]]. So at least internally we make a particular
# choice of typechecker.
#
# Deliberately not using @jaxtyped so that we share the same `memo` as whatever
# dynamic context we're currently in.
@typeguard.typechecked
def accepts_leaftype(x: cls.leaftype):
pass
if cls.leaftype is Any:
def is_leaftype(x, new_scope=True):
if new_scope and cls.structure is not None:
set_treepath_memo(None, cls.structure)
try:
accepts_leaftype(x)
except _TypeCheckError:
def is_flatten_leaftype(x):
return False
else:
return True
finally:
if new_scope and cls.structure is not None:
clear_treepath_memo()
leaves, structure = jtu.tree_flatten(obj, is_leaf=is_leaftype)
def is_check_leaftype(x, new_scope):
return True
else:
# We could use `isinstance` here but that would fail for more complicated
# types, e.g. PyTree[tuple[int]]. So at least internally we make a
# particular choice of typechecker.
#
# Deliberately not using @jaxtyped so that we share the same `memo` as
# whatever dynamic context we're currently in.
@typeguard.typechecked
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)
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)
if cls.structure is not None:
if cls.structure.isidentifier():
try:
@@ -119,10 +131,14 @@ class _MetaPyTree(type):
try:
prev_structure = pytree_memo[identifier]
except KeyError as e:
raise NameError(
f"Cannot process composite structure '{cls.structure}' as "
f"the structure name {identifier} has not been seen before."
) from e
jaxtyping_raise_from(
NameError(
f"Cannot process composite structure '{cls.structure}' "
f"as the structure name {identifier} has not been seen "
"before."
),
e,
)
# Not using `PyTreeDef.compose` due to JAX bug #18218.
prev_pytree = jtu.tree_unflatten(
prev_structure, [0] * prev_structure.num_leaves
@@ -152,7 +168,7 @@ class _MetaPyTree(type):
for leaf_index, leaf in enumerate(leaves):
if cls.structure is not None:
set_treepath_memo(leaf_index, cls.structure)
if not is_leaftype(leaf, new_scope=False):
if not is_check_leaftype(leaf, new_scope=False):
return False
clear_treepath_memo()
finally:
jaxtyping/_raise.py
+22
View File
@@ -0,0 +1,22 @@
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.
"""
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
+14 -8
View File
@@ -20,6 +20,8 @@
import threading
from typing import Optional
from ._raise import jaxtyping_raise
_shape_storage = threading.local()
@@ -70,10 +72,12 @@ 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:
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."
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."
)
)
if index is None:
_treepath_storage.value = f"~~delete~~({structure}) "
@@ -84,9 +88,11 @@ 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:
raise ValueError(
"Cannot use `?` annotations, e.g. `Shaped[Array, '?foo']`, except "
"when contained with structured `PyTree` annotations, e.g. "
"`PyTree[Shaped[Array, '?foo'], 'T']`."
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']`."
)
)
return _treepath_storage.value
test/conftest.py
+23
View File
@@ -41,6 +41,29 @@ def typecheck(request):
return request.param
@pytest.fixture(params=(False, True))
def jaxtyp(request):
import jaxtyping
if request.param:
# New-style
# @jaxtyping.jaxtyped(typechecker=typechecker)
# def f(...)
return lambda typechecker: jaxtyping.jaxtyped(typechecker=typechecker)
else:
# Old-style
# @jaxtyping.jaxtyped
# @typechecker
# def f(...)
def impl(typechecker):
def decorator(fn):
return jaxtyping.jaxtyped(typechecker(fn))
return decorator
return impl
@pytest.fixture()
def getkey():
def _getkey():
test/test_array.py
+49 -78
View File
@@ -33,7 +33,6 @@ from jaxtyping import (
Bool,
Float,
Float32,
jaxtyped,
PRNGKeyArray,
Shaped,
)
@@ -41,9 +40,8 @@ from jaxtyping import (
from .helpers import ParamError, ReturnError
def test_basic(typecheck):
@jaxtyped
@typecheck
def test_basic(jaxtyp, typecheck):
@jaxtyp(typecheck)
def g(x: Shaped[Array, "..."]):
pass
@@ -82,16 +80,14 @@ def test_dtypes():
assert key == val.__name__
def test_return(typecheck, getkey):
@jaxtyped
@typecheck
def test_return(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]:
return jnp.transpose(x)
g(jr.normal(getkey(), (3, 4)))
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def h(x: Float[Array, "b c"]) -> Float[Array, "b c"]:
return jnp.transpose(x)
@@ -99,9 +95,8 @@ def test_return(typecheck, getkey):
h(jr.normal(getkey(), (3, 4)))
def test_two_args(typecheck, getkey):
@jaxtyped
@typecheck
def test_two_args(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]):
return x @ y
@@ -109,8 +104,7 @@ def test_two_args(typecheck, getkey):
with pytest.raises(ParamError):
g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def h(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]) -> Shaped[Array, "b d"]:
return x @ y
@@ -119,9 +113,8 @@ def test_two_args(typecheck, getkey):
h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
def test_any_dtype(typecheck, getkey):
@jaxtyped
@typecheck
def test_any_dtype(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
return x
@@ -136,14 +129,12 @@ def test_any_dtype(typecheck, getkey):
g(jr.normal(getkey(), (1,)))
def test_nested_jaxtyped(typecheck, getkey):
@jaxtyped
@typecheck
def test_nested_jaxtyped(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float32[Array, "b c"], transpose: bool) -> Float32[Array, "c b"]:
return h(x, transpose)
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def h(x: Float32[Array, "c b"], transpose: bool) -> Float32[Array, "b c"]:
if transpose:
return jnp.transpose(x)
@@ -157,9 +148,8 @@ def test_nested_jaxtyped(typecheck, getkey):
g(jr.normal(getkey(), (2, 3)), False)
def test_nested_nojaxtyped(typecheck, getkey):
@jaxtyped
@typecheck
def test_nested_nojaxtyped(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float32[Array, "b c"]):
return h(x)
@@ -171,9 +161,8 @@ def test_nested_nojaxtyped(typecheck, getkey):
g(jr.normal(getkey(), (2, 3)))
def test_isinstance(typecheck, getkey):
@jaxtyped
@typecheck
def test_isinstance(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float32[Array, "b c"]) -> Float32[Array, " z"]:
y = jnp.transpose(x)
assert isinstance(y, Float32[Array, "c b"])
@@ -187,9 +176,8 @@ def test_isinstance(typecheck, getkey):
g(jr.normal(getkey(), (2, 3)))
def test_fixed(typecheck, getkey):
@jaxtyped
@typecheck
def test_fixed(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(
x: Float32[Array, "4 5 foo"], y: Float32[Array, " foo"]
) -> Float32[Array, "4 5"]:
@@ -204,9 +192,8 @@ def test_fixed(typecheck, getkey):
g(c, b)
def test_anonymous(typecheck, getkey):
@jaxtyped
@typecheck
def test_anonymous(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float32[Array, "foo _"], y: Float32[Array, " _"]):
pass
@@ -215,9 +202,8 @@ def test_anonymous(typecheck, getkey):
g(a, b)
def test_named_variadic(typecheck, getkey):
@jaxtyped
@typecheck
def test_named_variadic(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(
x: Float32[Array, "*batch foo"],
y: Float32[Array, " *batch"],
@@ -240,8 +226,7 @@ def test_named_variadic(typecheck, getkey):
with pytest.raises(ParamError):
g(a2, b1, c)
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def h(x: Float32[Array, " foo *batch"], y: Float32[Array, " foo *batch bar"]):
pass
@@ -255,9 +240,8 @@ def test_named_variadic(typecheck, getkey):
h(b, c)
def test_anonymous_variadic(typecheck, getkey):
@jaxtyped
@typecheck
def test_anonymous_variadic(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float32[Array, "... foo"], y: Float32[Array, " foo"]):
pass
@@ -277,9 +261,8 @@ def test_anonymous_variadic(typecheck, getkey):
g(a3, c)
def test_broadcast_fixed(typecheck, getkey):
@jaxtyped
@typecheck
def test_broadcast_fixed(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float32[Array, "#4"]):
pass
@@ -290,9 +273,8 @@ def test_broadcast_fixed(typecheck, getkey):
g(jr.normal(getkey(), (3,)))
def test_broadcast_named(typecheck, getkey):
@jaxtyped
@typecheck
def test_broadcast_named(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float32[Array, " #foo"], y: Float32[Array, " #foo"]):
pass
@@ -314,9 +296,8 @@ def test_broadcast_named(typecheck, getkey):
g(b, a)
def test_broadcast_variadic_named(typecheck, getkey):
@jaxtyped
@typecheck
def test_broadcast_variadic_named(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: Float32[Array, " *#foo"], y: Float32[Array, " *#foo"]):
pass
@@ -373,9 +354,8 @@ def test_broadcast_variadic_named(typecheck, getkey):
g(o, a)
def test_variadic_mixed_broadcast1(typecheck, getkey):
@jaxtyped
@typecheck
def test_variadic_mixed_broadcast(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def f(x: Float[Array, " *foo"], y: Float[Array, " #*foo"]):
pass
@@ -389,9 +369,8 @@ def test_variadic_mixed_broadcast1(typecheck, getkey):
f(c, d)
def test_variadic_mixed_broadcast2(typecheck, getkey):
@jaxtyped
@typecheck
def test_variadic_mixed_broadcast2(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def f(x: Float[Array, " *#foo"], y: Float[Array, " *foo"]):
pass
@@ -405,9 +384,8 @@ def test_variadic_mixed_broadcast2(typecheck, getkey):
f(c, d)
def test_variadic_mixed_broadcast3(typecheck, getkey):
@jaxtyped
@typecheck
def test_variadic_mixed_broadcast3(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def f(
x: Float[Array, "*B L D"],
*,
@@ -427,24 +405,20 @@ def test_no_commas():
Float32[Array, "foo, bar"]
def test_symbolic(typecheck, getkey):
@jaxtyped
@typecheck
def test_symbolic(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
return x[1:]
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
return jnp.concatenate([x, x])
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def bad_make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
return x
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def bad_cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
return jnp.concatenate([x, x, x])
@@ -464,9 +438,8 @@ def test_symbolic(typecheck, getkey):
bad_cat(x)
def test_incomplete_symbolic(typecheck, getkey):
@jaxtyped
@typecheck
def test_incomplete_symbolic(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def foo(x: Float32[Array, " 2*dim"]):
pass
@@ -573,9 +546,8 @@ def test_py310_unions():
assert isinstance(x, get_args(y))
def test_key(typecheck):
@jaxtyped
@typecheck
def test_key(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(x: PRNGKeyArray):
pass
@@ -592,7 +564,7 @@ def test_key(typecheck):
f(jnp.array(3.0))
def test_extension(typecheck, getkey):
def test_extension(jaxtyp, typecheck, getkey):
X = Shaped[Array, "a b"]
Y = Shaped[X, "c d"]
Z = Shaped[Array, "c d a b"]
@@ -601,8 +573,7 @@ def test_extension(typecheck, getkey):
X = Float[Array, "a"]
Y = Float[X, "b"]
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def f(a: X, b: Y):
...
test/test_decorator.py
+77
View File
@@ -1,9 +1,12 @@
import abc
import jax.random as jr
import pytest
from jaxtyping import Array, Float, jaxtyped
from .helpers import ParamError, ReturnError
class M(metaclass=abc.ABCMeta):
@jaxtyped
@@ -85,3 +88,77 @@ def test_context(getkey):
with jaxtyped("context"):
assert isinstance(a, Float[Array, "foo bar"])
assert not isinstance(b, Float[Array, "foo"])
def test_varargs(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(*args):
pass
f(1, 2)
def test_varkwargs(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(**kwargs):
pass
f(a=1, b=2)
def test_defaults(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(x, y=1):
pass
f(1)
class _GlobalFoo:
pass
def test_global_stringified_annotation(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(x: "_GlobalFoo") -> "_GlobalFoo":
return x
f(_GlobalFoo())
@jaxtyp(typecheck)
def g(x: int) -> "_GlobalFoo":
return x
@jaxtyp(typecheck)
def h(x: "_GlobalFoo") -> int:
return x
with pytest.raises(ReturnError):
g(1)
with pytest.raises(ParamError):
h(1)
# This test does not use `jaxtyp(typecheck)` because typeguard does some evil stack
# frame introspection to try and grab local variables.
def test_local_stringified_annotation(typecheck):
class LocalFoo:
pass
@jaxtyped(typechecker=typecheck)
def f(x: "LocalFoo") -> "LocalFoo":
return x
f(LocalFoo())
@jaxtyped
@typecheck
def g(x: "LocalFoo") -> "LocalFoo":
return x
g(LocalFoo())
# We don't check that errors are raised if it goes wrong, since we can't usually
# resolve local type annotations at runtime. Best we can hope for is not to raise
# a spurious error about not being able to find the type.
test/test_messages.py
+108
View File
@@ -0,0 +1,108 @@
from typing import Any
import equinox as eqx
import jax.numpy as jnp
import pytest
from jaxtyping import Array, Float, jaxtyped, PyTree, TypeCheckError
def test_arg_localisation(typecheck):
@jaxtyped(typechecker=typecheck)
def f(x: str, y: str, z: int):
pass
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: {}",
r"Parameter annotations: \(x: str, y: str, z: int\).",
]
for match in matches:
with pytest.raises(TypeCheckError, match=match):
f("hi", "bye", "not-an-int")
@jaxtyped(typechecker=typecheck)
def g(x: Float[Array, "a b"], y: Float[Array, "b c"]):
pass
x = jnp.zeros((2, 3))
y = jnp.zeros((4, 3))
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"Parameter annotations: \(x: Float\[Array, 'a b'\], y: "
r"Float\[Array, 'b c'\]\)."
),
"The current values for each jaxtyping axis annotation are as follows.",
"a=2",
"b=3",
]
for match in matches:
with pytest.raises(TypeCheckError, match=match):
g(x, y=y)
def test_return(typecheck):
@jaxtyped(typechecker=typecheck)
def f(x: PyTree[Any, " T"], y: PyTree[Any, " S"]) -> PyTree[Any, "T S"]:
return "foo"
x = (1, 2)
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}}",
"Return value: 'foo'",
r"Return annotation: PyTree\[Any, \"T S\"\].",
(
"The current values for each jaxtyping PyTree structure annotation are as "
"follows."
),
r"T=PyTreeDef\(\(\*, \*\)\)",
r"S=PyTreeDef\({'a': \*}\)",
]
for match in matches:
with pytest.raises(TypeCheckError, match=match):
f(x, y=y)
def test_dataclass_attribute(typecheck):
@jaxtyped(typechecker=typecheck)
class M(eqx.Module):
x: Float[Array, " *foo"]
y: PyTree[Any, " T"]
z: int
x = jnp.zeros((2, 3))
y = (1, (3, 4))
z = "not-an-int"
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"Parameter annotations: \(x: Float\[Array, '\*foo'\], "
r"y: PyTree\[Any, \"T\"\], z: int\)."
),
"The current values for each jaxtyping axis annotation are as follows.",
r"foo=\(2, 3\)",
(
"The current values for each jaxtyping PyTree structure annotation are as "
"follows."
),
r"T=PyTreeDef\(\(\*, \(\*, \*\)\)\)",
]
for match in matches:
with pytest.raises(TypeCheckError, match=match):
M(x, y, z)
test/test_pytree.py
+21 -31
View File
@@ -26,7 +26,7 @@ import jax.random as jr
import pytest
import jaxtyping
from jaxtyping import Array, Float, jaxtyped, PyTree
from jaxtyping import Array, Float, PyTree
from .helpers import make_mlp, ParamError
@@ -93,9 +93,8 @@ def test_nested_pytrees(getkey, typecheck):
g([1, 2, make_mlp()])
def test_pytree_array(typecheck):
@jaxtyped
@typecheck
def test_pytree_array(jaxtyp, typecheck):
@jaxtyp(typecheck)
def g(x: PyTree[Float[jnp.ndarray, "..."]]):
pass
@@ -107,9 +106,8 @@ def test_pytree_array(typecheck):
g(1.0)
def test_pytree_shaped_array(typecheck, getkey):
@jaxtyped
@typecheck
def test_pytree_shaped_array(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def g(x: PyTree[Float[jnp.ndarray, "b c"]]):
pass
@@ -196,9 +194,8 @@ def test_subclass_pytree():
assert not issubclass(int, PyTree)
def test_structure_match(typecheck):
@jaxtyped
@typecheck
def test_structure_match(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(x: PyTree[int, " T"], y: PyTree[str, " T"]):
pass
@@ -209,9 +206,8 @@ def test_structure_match(typecheck):
f(1, ("hi",))
def test_structure_prefix(typecheck):
@jaxtyped
@typecheck
def test_structure_prefix(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(x: PyTree[int, " T"], y: PyTree[str, "T ..."]):
pass
@@ -228,9 +224,8 @@ def test_structure_prefix(typecheck):
f((3, 4, 5), {"a": ("hi", "bye")})
def test_structure_suffix(typecheck):
@jaxtyped
@typecheck
def test_structure_suffix(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(x: PyTree[int, " T"], y: PyTree[str, "... T"]):
pass
@@ -245,9 +240,8 @@ def test_structure_suffix(typecheck):
f((3, 4, 5), {"a": ("hi", "bye")})
def test_structure_compose(typecheck):
@jaxtyped
@typecheck
def test_structure_compose(jaxtyp, typecheck):
@jaxtyp(typecheck)
def f(x: PyTree[int, " T"], y: PyTree[int, " S"], z: PyTree[str, "S T"]):
pass
@@ -262,8 +256,7 @@ def test_structure_compose(typecheck):
with pytest.raises(ParamError):
f((1, 2), {"a": 3}, ({"a": "hi"}, {"a": "bye"}))
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def g(x: PyTree[int, " T"], y: PyTree[int, " S"], z: PyTree[str, "T S"]):
pass
@@ -274,7 +267,7 @@ def test_structure_compose(typecheck):
@pytest.mark.parametrize("variadic", (False, True))
def test_treepath_dependence_function(variadic, typecheck, getkey):
def test_treepath_dependence_function(variadic, jaxtyp, typecheck, getkey):
if variadic:
jtshape = "*?foo"
shape = (2, 3)
@@ -282,8 +275,7 @@ def test_treepath_dependence_function(variadic, typecheck, getkey):
jtshape = "?foo"
shape = (4,)
@jaxtyped
@typecheck
@jaxtyp(typecheck)
def f(
x: PyTree[Float[Array, jtshape], " T"], y: PyTree[Float[Array, jtshape], " T"]
):
@@ -312,7 +304,7 @@ def test_treepath_dependence_dataclass(variadic, typecheck, getkey):
jtshape = "?foo"
shape = (4,)
@jaxtyping._decorator._jaxtyped_typechecker(typecheck)
@jaxtyping.jaxtyped(typechecker=typecheck)
class A(eqx.Module):
x: PyTree[Float[Array, jtshape], " T"]
y: PyTree[Float[Array, jtshape], " T"]
@@ -331,9 +323,8 @@ def test_treepath_dependence_dataclass(variadic, typecheck, getkey):
A((x1, x2), (y2, y1))
def test_treepath_dependence_missing_structure_annotation(typecheck, getkey):
@jaxtyped
@typecheck
def test_treepath_dependence_missing_structure_annotation(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def f(x: PyTree[Float[Array, "?foo"], " T"], y: PyTree[Float[Array, "?foo"]]):
pass
@@ -343,9 +334,8 @@ def test_treepath_dependence_missing_structure_annotation(typecheck, getkey):
f(x1, y1)
def test_treepath_dependence_multiple_structure_annotation(typecheck, getkey):
@jaxtyped
@typecheck
def test_treepath_dependence_multiple_structure_annotation(jaxtyp, typecheck, getkey):
@jaxtyp(typecheck)
def f(x: PyTree[PyTree[Float[Array, "?foo"], " S"], " T"]):
pass