Closures are being gc'd and causing failures to compile

Ch0ronomato: <p>Hi!</p>
<p>I’ve been contributing to pytensor adding the pytorch backend. The library generates a symbolic graph, and then</p>
<ol>
<li>Calls a series of specializations per operation</li>
<li>Glues them together with some generated code.</li>
</ol>
<p>An example of the generated code</p>
<pre data-code-wrap="py"><code class="lang-py">def pytorch_funcified_fgraph(x, y, z):
    # OpFromGraph{inline=False}(y, z)
    tensor_variable_2, tensor_variable_3 = pytorch_funcified_fgraph(y, z)
    # OpFromGraph{inline=False}(x, OpFromGraph{inline=False}.0)
    tensor_variable_5 = pytorch_funcified_fgraph_1(x, tensor_variable_2)
    # True_div(OpFromGraph{inline=False}.0, OpFromGraph{inline=False}.1)
    tensor_variable_6 = elemwise_fn_3(tensor_variable_5, tensor_variable_3)
    return (tensor_variable_6,)
</code></pre>
<p>and example of the dispatches</p>
<pre data-code-wrap="py"><code class="lang-py">@pytorch_funcify.register(Elemwise)
def pytorch_funcify_Elemwise(op, node, **kwargs):
    scalar_op = op.scalar_op
    base_fn = pytorch_funcify(scalar_op, node=node, **kwargs)

    def elemwise_fn(*inputs):
        Elemwise._check_runtime_broadcast(node, inputs)
        return base_fn(*inputs)

    return elemwise_fn
</code></pre>
<p>When all is said and done, you have a python function that glues all the dispatches together with some tensors. Torch wasn’t able to compile because it would fail to create a guard. The guard’s would fail to be created with this error</p>
<pre><code class="lang-auto">[2024-11-22 10:29:01,151] [44/0] torch._guards: [ERROR] Name: "G['__import_pytensor_dot_link_dot_utils'].elemwise_fn.__closure__[1].cell_contents"
[2024-11-22 10:29:01,151] [44/0] torch._guards: [ERROR]     Source: global
[2024-11-22 10:29:01,151] [44/0] torch._guards: [ERROR]     Create Function: TYPE_MATCH
[2024-11-22 10:29:01,151] [44/0] torch._guards: [ERROR]     Guard Types: None
[2024-11-22 10:29:01,151] [44/0] torch._guards: [ERROR]     Code List: None
[2024-11-22 10:29:01,151] [44/0] torch._guards: [ERROR]     Object Weakref: None
[2024-11-22 10:29:01,151] [44/0] torch._guards: [ERROR]     Guarded Class Weakref: None
</code></pre>
<p>The only way for us to get around this problem was either too disable the “leaf” nodes (in case of the error above, <code>elemwise_fn</code> would need to have a <code>torch.compiler.disable</code> decorator), run with <code>supress_errors = True</code>; or what I ended up doing in this PR, which is <a href="https://github.com/pymc-devs/pytensor/pull/1094" rel="noopener nofollow ugc">explicitly call setattr</a> on the <code>pytensor.link.utils</code> to be the return value of the <a href="https://github.com/pymc-devs/pytensor/pull/1094/files#diff-626f20cbb085bc9e19e127dc1a61c05de62a3628ef2372401886328b2e137525R48-R55" rel="noopener nofollow ugc">dispatch</a> function.</p>
<p>I did some digging through the compile logs, but I couldn’t quite pin down exactly what was happening at guard creation time. I could see that the generated code was doing a pretty good job at inlining all the necessary calls (it would look like this)</p>
<pre data-code-wrap="py"><code class="lang-py">class GraphModule(torch.nn.Module):
    def forward(self, L_x_ : torch.Tensor, L_y_ : torch.Tensor, L_z_ : torch.Tensor):
        input_4 = L_x_
        input_2 = L_y_
        input_3 = L_z_
        
        # File: /Users/ischweer/dev/pytensor/pytensor/link/pytorch/dispatch/elemwise.py:29, code: return base_fn(*inputs)
        tensor_variable = torch.add(input_4, input_2);  input_2 = None
        
        # File: /Users/ischweer/dev/pytensor/pytensor/link/pytorch/dispatch/elemwise.py:29, code: return base_fn(*inputs)
        tensor_variable_1 = torch.add(input_3, input_4);  input_3 = input_4 = None
        return (tensor_variable, tensor_variable_1)
</code></pre>
<p>But my best explanation for why the guards were failing is because of something being GC’d (but it could still be ran in the python runtime for…some reason), or pytorch had the module location for these methods. Is there any help I could get on narrowing down the issue?</p>
<p>I’m also attempting to make an MWE without pytensor to see if I can replicate this outside the library. It’s a little tricky though <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>CC <a class="mention" href="/u/ezyang">@ezyang</a> in case you would have an idea about the GC behavior.</p>