Torch.fx vs torch.jit.script()

username: mathmanu
cooked: <p>Hi,</p>
<p>I can see a new module in development called torch.fx with symbolic trace capability<br>
Could you please explain the difference between torch.fx and torch.jit.script()</p>
<p>Thanks &amp; Regards</p>

username: James_Reed
cooked: <p>Hello <a class="mention" href="/u/mathmanu">@mathmanu</a>,</p>
<p><code>torch.fx</code> is different from TorchScript in that it is a platform for Python-to-Python transformations of PyTorch code. TorchScript, on the other hand, is more targeted at moving PyTorch programs outside of Python for deployment purposes. In this sense, FX and TorchScript are orthogonal to each other, and can even be composed with each other (e.g. transform PyTorch programs with FX, then subsequently export to TorchScript for deployment).</p>
<p>Please stay tuned for more information about FX early next year. Note that FX is very unstable at this point and we do not recommend building off of the code in master at this time</p>

username: zetyquickly
cooked: <aside class="quote no-group" data-username="James_Reed" data-post="2" data-topic="100299">
<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/james_reed/48/34631_2.png" class="avatar"> James_Reed:</div>
<blockquote>
<p>Python-to-Python transformations</p>
</blockquote>
</aside>
<p>Could you please describe what does it mean?</p>
<blockquote>
<p><strong>Python code generation</strong> is what makes FX a Python-to-Python (or Module-to-Module) transformation toolkit. For each Graph IR, we can create valid Python code matching the Graph’s semantics. This functionality is wrapped up in <a href="https://pytorch.org/docs/stable/fx.html#torch.fx.GraphModule" rel="noopener nofollow ugc"><code>GraphModule</code></a>, which is a <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module" rel="noopener nofollow ugc"><code>torch.nn.Module</code></a> instance that holds a <a href="https://pytorch.org/docs/stable/fx.html#torch.fx.Graph" rel="noopener nofollow ugc"><code>Graph</code></a> as well as a <code>forward</code> method generated from the Graph.</p>
</blockquote>
<p>Why does one need to translate Python to Python? Is IR the same IR that we acquire during <code>jit.trace</code>?</p>

username: James_Reed
cooked: <p>Hi <a class="mention" href="/u/zetyquickly">@zetyquickly</a>,</p>
<blockquote>
<p>Could you please describe what does it mean?</p>
</blockquote>
<p>This concept is introduced in the documentation: <a href="https://pytorch.org/docs/stable/fx.html#module-torch.fx" class="inline-onebox" rel="noopener nofollow ugc">torch.fx — PyTorch 2.1 documentation</a>. FX produces valid Python <code>nn.Module</code> instances from its <code>Graph</code> representation</p>
<blockquote>
<p>Why does one need to translate Python to Python?</p>
</blockquote>
<p>FX emphasizes generating Python code so that it can be used within the existing PyTorch eager ecosystem. That is to say, code transformed by FX is not locked into one specific runtime (e.g. TorchScript) and all the normal tooling that can be used with normal PyTorch modules can be used with FX-generated modules.</p>
<blockquote>
<p>Is IR the same IR that we acquire during <code>jit.trace</code> ?</p>
</blockquote>
<p>No, the IR is not the same as that produced by <code>jit.trace</code>, intentionally so. FX is an entirely separate system that is superior to <code>jit.trace</code> in several ways:</p>
<ol>
<li>FX deeply integrates into the Python runtime, so it can better acquire accurate program representations, whereas <code>jit.trace</code> often silently captures wrong representations</li>
<li>FX’s Graph IR is much simpler since it represents the code at a higher level (<code>torch.nn</code> module calls are preserved rather than traced through). This is much easier to work with and understand and we have seen major productivity improvements using this IR</li>
</ol>