Fix dangling reference

Author: ischweer

pytensor/link/pytorch/dispatch/basic.py

@@ -61,20 +61,13 @@ def pytorch_funcify_FunctionGraph(
     conversion_func=pytorch_funcify,
     **kwargs,
 ):
-    def constants_wrapper(x, **kwargs):
-        x = pytorch_typify(x)
-
-        @torch.compiler.assume_constant_result
-        def torch_assume_constant(arg=x):
-            return arg
-
-        return torch_assume_constant
+    if "type_conversion_fn" not in kwargs:
+        kwargs["type_conversion_fn"] = pytorch_typify
 
     built_kwargs = {"conversion_func": conversion_func, **kwargs}
     return fgraph_to_python(
         fgraph,
         conversion_func,
-        type_conversion_fn=constants_wrapper,
         fgraph_name=fgraph_name,
         **built_kwargs,
     )

pytensor/link/pytorch/linker.py

@@ -10,7 +10,9 @@ def __init__(self, *args, **kwargs):
         self.gen_functors = []
 
     def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
-        from pytensor.link.pytorch.dispatch import pytorch_funcify
+        import torch
+
+        from pytensor.link.pytorch.dispatch import pytorch_funcify, pytorch_typify
 
         # We want to have globally unique names
         # across the entire pytensor graph, not
@@ -25,9 +27,21 @@ def conversion_func_register(*args, **kwargs):
             self.gen_functors.append((f"_{name}", functor))
             return functor
 
+        def constants_wrapper(x, **kwargs):
+            x = pytorch_typify(x)
+
+            @torch.compiler.assume_constant_result
+            def torch_assume_constant(arg=x):
+                return arg
+
+            name = kwargs["unique_name"](torch_assume_constant)
+            self.gen_functors.append((f"_{name}", torch_assume_constant))
+            return torch_assume_constant
+
         built_kwargs = {
             "unique_name": generator,
             "conversion_func": conversion_func_register,
+            "type_conversion_fn": constants_wrapper,
             **kwargs,
         }
         return pytorch_funcify(

pytensor/link/utils.py

@@ -766,8 +766,8 @@ def fgraph_to_python(
                     global_env[local_input_name] = type_conversion_fn(
                         input_storage[0], variable=i, storage=input_storage, **kwargs
                     )
-                # TODO: We could attempt to use the storage arrays directly
-                # E.g. `local_input_name = f"{local_input_name}[0]"`
+            # TODO: We could attempt to use the storage arrays directly
+            # E.g. `local_input_name = f"{local_input_name}[0]"`
             node_input_names.append(local_input_name)
 
         node_output_names = [unique_name(v) for v in node.outputs]