Share one name generator

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -1,6 +1,5 @@
 import inspect
 import sys
-from collections import defaultdict
 from functools import singledispatch
 from types import NoneType
 
@@ -59,19 +58,12 @@ 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__
-            if fnames_counter[name]:
-                name = f"{name}_{fnames_counter[name]}"
-            print("Registering functor", name)
-            setattr(module, name, functor)
-            fnames_counter[functor.__name__] += 1
+            setattr(module, kwargs["unique_name"](functor), functor)
         return functor
 
     return fgraph_to_python(
@@ -191,7 +183,6 @@ def pytorch_funcify_OpFromGraph(op, node, **kwargs):
 
     # Apply inner rewrites
     PYTORCH.optimizer(op.fgraph)
-
     fgraph_fn = pytorch_funcify(op.fgraph, **kwargs, squeeze_output=True)
     return fgraph_fn
 

pytensor/link/pytorch/linker.py

@@ -2,6 +2,7 @@
 
 from pytensor.graph.basic import Variable
 from pytensor.link.basic import JITLinker
+from pytensor.link.utils import unique_name_generator
 
 
 class PytorchLinker(JITLinker):
@@ -18,8 +19,10 @@ def output_filter(self, var: Variable, out: Any) -> Any:
     def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
         from pytensor.link.pytorch.dispatch import pytorch_funcify
 
+        generator = unique_name_generator(["torch_linker"])
+        built_kwargs = {"unique_name": generator, **kwargs}
         return pytorch_funcify(
-            fgraph, input_storage=input_storage, storage_map=storage_map, **kwargs
+            fgraph, input_storage=input_storage, storage_map=storage_map, **built_kwargs
         )
 
     def jit_compile(self, fn):

pytensor/link/utils.py

@@ -704,6 +704,8 @@ def fgraph_to_python(
     get_name_for_object
         A function used to provide names for the objects referenced within the
         generated function.
+    unique_name
+        A function to make random function names for generated code (access through kwargs)
     squeeze_output
         If the `FunctionGraph` has only one output and this option is
         ``True``, return the single output instead of a tuple with the output.
@@ -717,7 +719,8 @@ def fgraph_to_python(
     if storage_map is None:
         storage_map = {}
 
-    unique_name = unique_name_generator([fgraph_name])
+    unique_name = kwargs.get("unique_name", unique_name_generator([fgraph_name]))
+    kwargs["unique_name"] = unique_name
 
     if global_env is None:
         global_env = {}