PR comments

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -54,22 +54,16 @@ def pytorch_funcify_FunctionGraph(
     fgraph,
     node=None,
     fgraph_name="pytorch_funcified_fgraph",
+    conversion_func=pytorch_funcify,
     **kwargs,
 ):
-    # Ensure that torch is aware of the leaf nodes
-    # of generated code so we can compile code
-    # without graph breaks
-    def conversion_func_register(*args, **kwargs):
-        functor = pytorch_funcify(*args, **kwargs)
-        kwargs["linker"].record_fn(kwargs["unique_name"](functor), functor)
-        return functor
-
+    built_kwargs = {"conversion_func": conversion_func, **kwargs}
     return fgraph_to_python(
         fgraph,
-        conversion_func_register,
+        conversion_func,
         type_conversion_fn=pytorch_typify,
         fgraph_name=fgraph_name,
-        **kwargs,
+        **built_kwargs,
     )
 
 

pytensor/link/pytorch/linker.py

@@ -28,7 +28,20 @@ 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, "linker": self, **kwargs}
+
+        # Ensure that torch is aware of the generated
+        # code so we can compile without graph breaks
+        def conversion_func_register(*args, **kwargs):
+            functor = pytorch_funcify(*args, **kwargs)
+            name = kwargs["unique_name"](functor)
+            self.gen_functors.append((f"_{name}", functor))
+            return functor
+
+        built_kwargs = {
+            "unique_name": generator,
+            "conversion_func": conversion_func_register,
+            **kwargs,
+        }
         return pytorch_funcify(
             fgraph, input_storage=input_storage, storage_map=storage_map, **built_kwargs
         )
@@ -37,6 +50,16 @@ def jit_compile(self, fn):
         import torch
 
         class wrapper:
+            """
+            Pytorch would fail compiling our method when trying
+            to resolve some of the methods returned from dispatch
+            calls. We want to be careful to not leak the methods,
+            so this class just holds them and provisions the expected
+            location accordingly
+
+            https://discuss.pytorch.org/t/closures-are-being-gcd-and-causing-failures-to-compile/213319
+            """
+
             def __init__(self, fn, gen_functors):
                 self.fn = torch.compile(fn)
                 self.gen_functors = copy.copy(gen_functors)
@@ -46,18 +69,18 @@ def __call__(self, *args, **kwargs):
 
                 # set attrs
                 for n, fn in self.gen_functors:
-                    setattr(pytensor.link.utils, n, fn)
+                    setattr(pytensor.link.utils, n[1:], fn)
 
                 res = self.fn(*args, **kwargs)
 
                 # unset attrs
                 for n, _ in self.gen_functors:
-                    delattr(pytensor.link.utils, n)
+                    if getattr(pytensor.link.utils, n, False):
+                        delattr(pytensor.link.utils, n[1:])
 
                 return res
 
             def __del__(self):
-                print("del")
                 del self.gen_functors
 
         res = wrapper(fn, self.gen_functors)
@@ -71,6 +94,3 @@ 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))

pytensor/link/utils.py

@@ -673,6 +673,7 @@ def fgraph_to_python(
     local_env: dict[Any, Any] | None = None,
     get_name_for_object: Callable[[Any], str] = get_name_for_object,
     squeeze_output: bool = False,
+    unique_name: Callable | None = None,
     **kwargs,
 ) -> Callable:
     """Convert a `FunctionGraph` into a regular Python function.
@@ -705,7 +706,7 @@ def fgraph_to_python(
         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)
+        A function to make random function names for generated code
     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.
@@ -719,7 +720,10 @@ def fgraph_to_python(
     if storage_map is None:
         storage_map = {}
 
-    unique_name = kwargs.get("unique_name", unique_name_generator([fgraph_name]))
+    if not unique_name:
+        unique_name = unique_name_generator([fgraph_name])
+
+    # make sure we plumb this through
     kwargs["unique_name"] = unique_name
 
     if global_env is None: