diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py index a269d00..bb1ac6d 100644 --- a/jaxtyping/decorator.py +++ b/jaxtyping/decorator.py @@ -21,17 +21,24 @@ import dataclasses import functools as ft import inspect import threading +import weakref storage = threading.local() +_fns = weakref.WeakKeyDictionary() + + class _Jaxtyped: def __init__(self, fn): - self.fn = fn + # Stored externally so that it doesn't get blatted in the `ft.wraps` below by + # a function that already has a `fn` attribute. + _fns[self] = fn def __get__(self, instance, owner): - return ft.wraps(self.fn)(_Jaxtyped(self.fn.__get__(instance, owner))) + fn = _fns[self] + return ft.wraps(fn)(_Jaxtyped(fn.__get__(instance, owner))) def __call__(self, *args, **kwargs): try: @@ -39,8 +46,9 @@ class _Jaxtyped: except AttributeError: memo_stack = storage.memo_stack = [] memo_stack.append(({}, {}, {})) + fn = _fns[self] try: - return self.fn(*args, **kwargs) + return fn(*args, **kwargs) finally: memo_stack.pop() diff --git a/test/test_threading.py b/test/test_threading.py index 557fba7..6bb6210 100644 --- a/test/test_threading.py +++ b/test/test_threading.py @@ -31,8 +31,6 @@ class _ErrorableThread(threading.Thread): super().run() except Exception as e: self.exc = e - finally: - del self._target, self._args, self._kwargs def join(self, timeout=None): super().join(timeout)