From e5cc75e4a309fd0e41a4e5ddb74f5805a31b5656 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sat, 2 Dec 2023 17:30:30 -0800
Subject: [PATCH] Removed internal `jaxtyped_fns` registry that is no longer
 needed.

---
 jaxtyping/_decorator.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
index fc7a6f7..a36fb88 100644
--- a/jaxtyping/_decorator.py
+++ b/jaxtyping/_decorator.py
@@ -22,9 +22,7 @@ import functools as ft
 import inspect
 import itertools as it
 import sys
-import types
 import warnings
-import weakref
 from typing import Any, get_args, get_origin, get_type_hints, overload
 
 
@@ -40,9 +38,6 @@ from ._config import config
 from ._storage import pop_shape_memo, push_shape_memo
 
 
-_jaxtyped_fns = weakref.WeakSet()
-
-
 class TypeCheckError(TypeError):
     pass
 
@@ -244,8 +239,6 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
 
     if fn is _sentinel:
         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
@@ -472,7 +465,6 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
                 finally:
                     pop_shape_memo()
 
-        _jaxtyped_fns.add(wrapped_fn)
         return wrapped_fn