Use explicit closure lifetime and tracking

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -5,6 +5,7 @@
 import torch
 import torch.compiler
 
+import pytensor.link.utils
 from pytensor.compile import PYTORCH
 from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.ops import DeepCopyOp
@@ -56,9 +57,15 @@ def pytorch_funcify_FunctionGraph(
     fgraph_name="pytorch_funcified_fgraph",
     **kwargs,
 ):
+    def conversion_func_register(*args, **kwargs):
+        functor = pytorch_funcify(*args, **kwargs)
+        module = pytensor.link.utils
+        setattr(module, kwargs["unique_name"](functor), functor)
+        return functor
+
     return fgraph_to_python(
         fgraph,
-        pytorch_funcify,
+        conversion_func_register,
         type_conversion_fn=pytorch_typify,
         fgraph_name=fgraph_name,
         **kwargs,
@@ -173,7 +180,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)
     # Disable one step inlining to prevent torch from trying to import local functions
     # defined in `pytorch_funcify`

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 = {}

tests/link/pytorch/test_basic.py

@@ -335,7 +335,7 @@ def test_pytorch_OpFromGraph():
     ofg_2 = OpFromGraph([x, y], [x * y, x - y])
 
     o1, o2 = ofg_2(y, z)
-    out = ofg_1(x, o1) + o2
+    out = ofg_1(x, o1) / o2
 
     xv = np.ones((2, 2), dtype=config.floatX)
     yv = np.ones((2, 2), dtype=config.floatX) * 3