Make module references FIFO

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -1,5 +1,6 @@
 import inspect
 import sys
+from collections import defaultdict
 from functools import singledispatch
 from types import NoneType
 
@@ -58,23 +59,19 @@ def pytorch_funcify_FunctionGraph(
     fgraph_name="pytorch_funcified_fgraph",
     **kwargs,
 ):
+    fnames_counter = defaultdict(int)
+
     def conversion_func_register(*args, **kwargs):
         last_frame = inspect.currentframe().f_back
         functor = pytorch_funcify(*args, **kwargs)
         if last_frame:
             module = sys.modules[last_frame.f_globals["__name__"]]
             name = functor.__name__
-            i = 0
-            prefix = ""
-            while True:
-                name = functor.__name__ + prefix
-                if name in dir(module):
-                    i += 1
-                    prefix = f"_{i}"
-                else:
-                    print("Registering functor", name)
-                    setattr(module, name, functor)
-                    break
+            if fnames_counter[name]:
+                name = f"{name}_{fnames_counter[name]}"
+            print("Registering functor", name)
+            setattr(module, name, functor)
+            fnames_counter[functor.__name__] += 1
         return functor
 
     return fgraph_to_python(