Make lifetime of closure explicit

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -1,3 +1,5 @@
+import inspect
+import sys
 from functools import singledispatch
 from types import NoneType
 
@@ -56,9 +58,28 @@ def pytorch_funcify_FunctionGraph(
     fgraph_name="pytorch_funcified_fgraph",
     **kwargs,
 ):
+    def conversion_func_register(*args, **kwargs):
+        last_frame = inspect.currentframe().f_back
+        functor = pytorch_funcify(*args, **kwargs)
+        if last_frame:
+            module = sys.modules[last_frame.f_globals["__name__"]]
+            name = functor.__name__
+            i = 0
+            prefix = ""
+            while True:
+                name = functor.__name__ + prefix
+                if name in dir(module):
+                    i += 1
+                    prefix = f"_{i}"
+                else:
+                    print("Registering functor", name)
+                    setattr(module, name, functor)
+                    break
+        return functor
+
     return fgraph_to_python(
         fgraph,
-        pytorch_funcify,
+        conversion_func_register,
         type_conversion_fn=pytorch_typify,
         fgraph_name=fgraph_name,
         **kwargs,
@@ -175,9 +196,7 @@ def pytorch_funcify_OpFromGraph(op, node, **kwargs):
     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`
-    return torch.compiler.disable(fgraph_fn, recursive=False)
+    return fgraph_fn
 
 
 @pytorch_funcify.register(TensorFromScalar)

pytensor/link/pytorch/dispatch/blockwise.py

@@ -16,7 +16,6 @@ def funcify_Blockwise(op: Blockwise, node, *args, **kwargs):
     for _ in range(batched_dims):
         inner_func = torch.vmap(inner_func)
 
-    @torch.compiler.disable(recursive=False)
     def batcher(*inputs):
         op._check_runtime_broadcast(node, inputs)
         # broadcast on batched_dims

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

tests/link/pytorch/test_blockwise.py

@@ -29,7 +29,6 @@ def perform(self, *_):
 
 @basic.pytorch_funcify.register(TestOp)
 def evaluate_test_op(op, **_):
-    @torch.compiler.disable(recursive=False)
     def func(a, b):
         op.call_shapes.extend(map(torch.Tensor.size, [a, b]))
         return a @ b