diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 2c3fd73..c6cc612 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -19,6 +19,7 @@ import importlib.metadata import typing +import warnings # First import some things as normal from ._array_types import ( @@ -196,4 +197,20 @@ elif has_jax: del has_jax +check_equinox_version = True # easy-to-replace line with copybara +if check_equinox_version: + try: + eqx_version = importlib.metadata.version("equinox") + except importlib.metadata.PackageNotFoundError: + pass + else: + major, minor, patch = eqx_version.split(".") + equinox_version = (int(major), int(minor), int(patch)) + if equinox_version < (0, 11, 0): + warnings.warn( + "jaxtyping version >=0.2.23 should be used with Equinox version " + ">=0.11.1" + ) + + __version__ = importlib.metadata.version("jaxtyping") diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py index a0c4f13..40a35d7 100644 --- a/jaxtyping/_decorator.py +++ b/jaxtyping/_decorator.py @@ -23,6 +23,7 @@ import inspect import threading import types import weakref +from typing import get_args, get_origin try: @@ -72,7 +73,7 @@ def jaxtyped(fn): 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 axes sizes as the function it was called from. + that this function use the same axis sizes as the function it was called from. Likewise, this means you can use `isinstance` checks inside a function body and have them contribute to the same collection of consistency checks performed @@ -134,7 +135,59 @@ def jaxtyped(fn): return wrapped_fn +@jaxtyped +def _check_dataclass_annotations(self, typechecker): + 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 + if isinstance(annotation, str): + # Don't support 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) + if len(args) == 1 and isinstance(args[0], str): + # We also special-case this one kind of partially-stringified type + # annotation, so as to support Equinox