diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 0250b00..c0b270f 100644 --- a/jaxtyping/__init__.py +++ b/jaxtyping/__init__.py @@ -53,4 +53,4 @@ from .import_hook import install_import_hook from .pytree_type import PyTree -__version__ = "0.2.1" +__version__ = "0.2.2" diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py index ae5874e..c03a09a 100644 --- a/jaxtyping/array_types.py +++ b/jaxtyping/array_types.py @@ -388,7 +388,7 @@ class _MetaAbstractDtype(type): name = "Array" else: raise ValueError(f"array_name_format {_array_name_format} not recognised") - return _MetaAbstractArray( + out = _MetaAbstractArray( name, (AbstractArray,), dict( @@ -398,6 +398,8 @@ class _MetaAbstractDtype(type): index_variadic=index_variadic, ), ) + out.__module__ = "jaxtyping" + return out class AbstractDtype(metaclass=_MetaAbstractDtype): @@ -468,6 +470,7 @@ else: _Cls.__name__ = name _Cls.__qualname__ = name + _Cls.__module__ = "jaxtyping" return _Cls UInt8 = _make_dtype(_uint8, "UInt8") diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py index 68ceb0c..b4916f5 100644 --- a/jaxtyping/decorator.py +++ b/jaxtyping/decorator.py @@ -24,17 +24,24 @@ import threading storage = threading.local() -def jaxtyped(fn): - @ft.wraps(fn) - def wrapper(*args, **kwargs): +class _Jaxtyped: + def __init__(self, fn): + self.fn = fn + + def __get__(self, instance, owner): + return ft.wraps(self.fn)(_Jaxtyped(self.fn.__get__(instance, owner))) + + def __call__(self, *args, **kwargs): try: memo_stack = storage.memo_stack except AttributeError: memo_stack = storage.memo_stack = [] memo_stack.append(({}, {}, {})) try: - return fn(*args, **kwargs) + return self.fn(*args, **kwargs) finally: memo_stack.pop() - return wrapper + +def jaxtyped(fn): + return ft.wraps(fn)(_Jaxtyped(fn)) diff --git a/jaxtyping/pytree_type.py b/jaxtyping/pytree_type.py index c19a0bc..38a7008 100644 --- a/jaxtyping/pytree_type.py +++ b/jaxtyping/pytree_type.py @@ -33,6 +33,7 @@ class _FakePyTree(Generic[_T]): _FakePyTree.__name__ = "PyTree" _FakePyTree.__qualname__ = "PyTree" +_FakePyTree.__module__ = "builtins" # Can't do type("PyTree", (Generic[_T],), {}) because dynamic subclassing of typeforms # isn't allowed. # Can't do types.new_class("PyTree", (Generic[_T],), {}) because that has __module__ @@ -49,7 +50,9 @@ class _MetaPyTree(type): @ft.lru_cache(maxsize=None) def __getitem__(cls, item): name = str(_FakePyTree[item]) - return _MetaSubscriptPyTree(name, (), {"leaftype": item}) + out = _MetaSubscriptPyTree(name, (), {"leaftype": item}) + out.__module__ = "jaxtyping" + return out class _MetaSubscriptPyTree(type): @@ -80,6 +83,7 @@ class _MetaSubscriptPyTree(type): PyTree = _MetaPyTree("PyTree", (), {}) +PyTree.__module__ = "jaxtyping" # Can't do `class PyTree(Generic[_T]): ...` because we need to override the # instancecheck for PyTree[foo], but we subclassing # `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed. diff --git a/test/test_decorator.py b/test/test_decorator.py new file mode 100644 index 0000000..696b25c --- /dev/null +++ b/test/test_decorator.py @@ -0,0 +1,13 @@ +from jaxtyping import jaxtyped + + +class M: + @jaxtyped + @classmethod + def f(cls): + return 3 + + +# Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod +def test_decorator(): + assert M.f() == 3