Clean up after ourselves

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -5,7 +5,6 @@
 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
@@ -62,8 +61,7 @@ def pytorch_funcify_FunctionGraph(
     # without graph breaks
     def conversion_func_register(*args, **kwargs):
         functor = pytorch_funcify(*args, **kwargs)
-        module = pytensor.link.utils
-        setattr(module, kwargs["unique_name"](functor), functor)
+        kwargs["linker"].record_fn(kwargs["unique_name"](functor), functor)
         return functor
 
     return fgraph_to_python(

pytensor/link/pytorch/linker.py

@@ -1,3 +1,4 @@
+import copy
 from typing import Any
 
 from pytensor.graph.basic import Variable
@@ -8,6 +9,10 @@
 class PytorchLinker(JITLinker):
     """A `Linker` that compiles NumPy-based operations using torch.compile."""
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gen_functors = []
+
     def input_filter(self, inp: Any) -> Any:
         from pytensor.link.pytorch.dispatch import pytorch_typify
 
@@ -23,15 +28,41 @@ def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
         # across the entire pytensor graph, not
         # just the subgraph
         generator = unique_name_generator(["torch_linker"])
-        built_kwargs = {"unique_name": generator, **kwargs}
+        built_kwargs = {"unique_name": generator, "linker": self, **kwargs}
         return pytorch_funcify(
             fgraph, input_storage=input_storage, storage_map=storage_map, **built_kwargs
         )
 
     def jit_compile(self, fn):
         import torch
 
-        return torch.compile(fn)
+        class wrapper:
+            def __init__(self, fn, gen_functors):
+                self.fn = torch.compile(fn)
+                self.gen_functors = copy.copy(gen_functors)
+
+            def __call__(self, *args, **kwargs):
+                import pytensor.link.utils
+
+                # set attrs
+                for n, fn in self.gen_functors:
+                    setattr(pytensor.link.utils, n, fn)
+
+                res = self.fn(*args, **kwargs)
+
+                # unset attrs
+                for n, _ in self.gen_functors:
+                    delattr(pytensor.link.utils, n)
+
+                return res
+
+            def __del__(self):
+                print("del")
+                del self.gen_functors
+
+        res = wrapper(fn, self.gen_functors)
+        self.gen_functors = []
+        return res
 
     def create_thunk_inputs(self, storage_map):
         thunk_inputs = []
@@ -40,3 +71,6 @@ def create_thunk_inputs(self, storage_map):
             thunk_inputs.append(sinput)
 
         return thunk_inputs
+
+    def record_fn(self, name, fn):
+        self.gen_functors.append((name, fn))