diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index cf9441c..932a5f9 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -67,4 +67,4 @@ from .import_hook import install_import_hook from .pytree_type import PyTree -__version__ = "0.2.7" +__version__ = "0.2.8" diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py index b4916f5..61d82bb 100644 --- a/jaxtyping/decorator.py +++ b/jaxtyping/decorator.py @@ -18,6 +18,7 @@ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. import functools as ft +import inspect import threading @@ -44,4 +45,9 @@ class _Jaxtyped: def jaxtyped(fn): - return ft.wraps(fn)(_Jaxtyped(fn)) + if inspect.isclass(fn): # allow decorators on class definitions + init = jaxtyped(fn.__init__) + fn.__init__ = init + return fn + else: + return ft.wraps(fn)(_Jaxtyped(fn)) diff --git a/setup.py b/setup.py index 7185d91..6541801 100644 --- a/setup.py +++ b/setup.py @@ -67,11 +67,14 @@ python_requires = "~=3.7" # We use typeguard internally (in a fairly minimal way), but it's not required that # end users make the same choice. +# For typing_extensions, we choose versions that match +# https://github.com/explosion/confection/blob/main/setup.cfg#L33 used in colab + install_requires = [ "jax>=0.3.4", "numpy>=1.20.0", "typeguard>=2.13.3", - "typing_extensions>=4.2.0", + "typing_extensions>=3.7.4.1", ] entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"])