PR comments

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -56,12 +56,15 @@ def pytorch_funcify_FunctionGraph(
     fgraph_name="pytorch_funcified_fgraph",
     **kwargs,
 ):
-    # Ensure that torch is aware of the leaf nodes
-    # of generated code so we can compile code
-    # without graph breaks
+    # 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)
-        kwargs["linker"].record_fn(kwargs["unique_name"](functor), functor)
+
+        # The linker could be missing cause we're only using
+        # function graph
+        if {"linker", "unique_name"}.issubset(set(kwargs.keys())):
+            kwargs["linker"].record_fn(kwargs["unique_name"](functor), functor)
         return functor
 
     return fgraph_to_python(

pytensor/link/pytorch/linker.py

@@ -37,6 +37,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 +56,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)
@@ -73,4 +83,4 @@ def create_thunk_inputs(self, storage_map):
         return thunk_inputs
 
     def record_fn(self, name, fn):
-        self.gen_functors.append((name, fn))
+        self.gen_functors.append((f"_{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: