Subclass of `fx.GraphModule` generating its own code

Linux-cpp-lisp: <p>Hi all,</p>
<p>Is it possible to neatly encapsulate a <code>torch.fx.GraphModule</code> that generates its own graph? (Based, say, on <code>__init__</code> parameters.)</p>
<p>The obvious approach appears not to work:</p>
<pre><code class="lang-python">In [8]: class Mod(fx.GraphModule):
   ...:     def __init__(self):
   ...:         super().__init__(self, fx.Graph())
   ...:         def thingy(x):
   ...:             return x* 2.
   ...:         self.graph = fx.symbolic_trace(thingy)
   ...:         self.recompile()
   ...: 

In [9]: m = Mod()

In [10]: m(torch.ones(4))
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
&lt;ipython-input-10-56e7b4f5b37c&gt; in &lt;module&gt;
----&gt; 1 m(torch.ones(4))

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/fx/graph_module.py in wrapped_call(self, *args, **kwargs)
    306             try:
    307                 sys.excepthook = print_full_traceback
--&gt; 308                 return cls_call(self, *args, **kwargs)
    309             finally:
    310                 sys.excepthook = old_excepthook

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/fx/graph_module.py in wrapped_call(self, *args, **kwargs)
    306             try:
    307                 sys.excepthook = print_full_traceback
--&gt; 308                 return cls_call(self, *args, **kwargs)
    309             finally:
    310                 sys.excepthook = old_excepthook

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--&gt; 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

TypeError: forward() takes 1 positional argument but 2 were given
</code></pre>
<p>Adding an explicit <code>self</code> breaks <code>recompile()</code>:</p>
<pre><code class="lang-auto">In [11]: class Mod(fx.GraphModule):
    ...:     def __init__(self):
    ...:         super().__init__(self, fx.Graph())
    ...:         def thingy(self, x):
    ...:             return x* 2.
    ...:         self.graph = fx.symbolic_trace(thingy)
    ...:         self.recompile()
    ...: 

In [12]: m = Mod()
Traceback (most recent call last):
...
    self.recompile()

  File "&lt;eval_with_key_9&gt;", line 2
    def forward(self, self, x):
    ^
SyntaxError: duplicate argument 'self' in function definition
</code></pre>
<p>Is this sort of architecture — module that generates its own compute graph — supported?</p>
<p>Thanks!</p>

Linux-cpp-lisp: <p>Well nevermind… this was a small mistake but I’m leaving the post in case it happens to anyone else:</p>
<pre><code class="lang-python">In [17]: class Mod(fx.GraphModule):
    ...:     def __init__(self):
    ...:         super().__init__(self, fx.Graph())
    ...:         def thingy(x):
    ...:             return x* 2.
    ...:         self.graph = fx.symbolic_trace(thingy).graph
    ...:         self.recompile()
    ...: 

In [18]: m = Mod()

In [19]: m(torch.ones(4))
Out[19]: tensor([2., 2., 2., 2.])
</code></pre>
<p>It works when I set the <code>self.graph</code> to <code>symbolic_trace(...).graph</code>.</p>
<p>This raises a different question: why doesn’t it throw an error when I try to assign a <code>GraphModule</code> to <code>self.graph</code>?</p>

James_Reed: <aside class="quote no-group" data-username="Linux-cpp-lisp" data-post="2" data-topic="115140">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/linux-cpp-lisp/48/34184_2.png" class="avatar"> Linux-cpp-lisp:</div>
<blockquote>
<p>This raises a different question: why doesn’t it throw an error when I try to assign a <code>GraphModule</code> to <code>self.graph</code>?</p>
</blockquote>
</aside>
<blockquote>
<p>This raises a different question: why doesn’t it throw an error when I try to assign a <code>GraphModule</code> to <code>self.graph</code>?</p>
</blockquote>
<p>This is a good question. We rely on Python type annotations to ensure type correctness. Typically in our development we use <code>mypy</code> to check type-correctness of code. However, in this case it seems that <code>graph</code>’s setter doesn’t have an annotation for its parameter. I will go ahead and add one as well as an assert on an isinstance check. Thanks for bringing this up!</p>

Linux-cpp-lisp: <p>Makes sense, thanks!</p>
<p><a class="mention" href="/u/james_reed">@James_Reed</a> I do have a follow up question on this that is a little less trivial. Since its the same subject as the original thread title, I figured I’d post it here:</p>
<p>How can I generate a graph for <code>forward()</code> of a <code>GraphModule</code> that has access to self? For example:</p>
<pre><code class="lang-python">In [1]: import torch

In [2]: from torch import fx

In [3]: class M(fx.GraphModule):
   ...:     def __init__(self):
   ...:         super().__init__(self, fx.Graph())
   ...:         graph = fx.symbolic_trace(M.fake_forward)
   ...:         self.graph = graph
   ...:         self.recompile()
   ...:         self.register_buffer("w1", torch.ones(3))
   ...:     def fake_forward(self, x):
   ...:         return self.w1 * x
   ...: 

In [4]: m = M()
</code></pre>
<p>Gives:</p>
<pre><code class="lang-auto">  File "&lt;eval_with_key_1&gt;", line 2
    def forward(self, self, x):
    ^
SyntaxError: duplicate argument 'self' in function definition
</code></pre>
<p>Is this possible?</p>

James_Reed: <aside class="quote no-group" data-username="Linux-cpp-lisp" data-post="4" data-topic="115140">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/linux-cpp-lisp/48/34184_2.png" class="avatar"> Linux-cpp-lisp:</div>
<blockquote>
<p><code>recompile</code></p>
</blockquote>
</aside>
<p>Hi <a class="mention" href="/u/linux-cpp-lisp">@Linux-cpp-lisp</a>,</p>
<p>The reason this is happening is because if the argument to <code>symbolic_trace</code> is not an <code>nn.Module</code>, we treat it as a free function, regardless of whether it’s a class method or not. I’ll file an issue and see if we can figure out a way to distinguish this case v.s. the truly-free-function case: <a href="https://github.com/pytorch/pytorch/issues/54785" class="inline-onebox" rel="noopener nofollow ugc">[FX] Cannot symbolically trace class method · Issue #54785 · pytorch/pytorch · GitHub</a></p>