Alternative to TorchScript that supports control flow and autograd

IgnacioPickering: <p>I’m currently relying on torchscript for many models, since I need to execute them from C++. I recently learned however that there are no plans to keep developing JIT.</p>
<p>From the documentation I gather the alternatives are <code>torch.compile</code>, if running in python, and <code>torch.export</code> if running in a python-less environment.</p>
<p>I believe that neither of these are fit for my use case, since I need to call autograd for inference, and my model requieres some control flow. Even if worked to fully get rid of the control flow (unlikely to be possible) the lack of autograd is a dealbreaker.</p>
<p>I wanted to know what are my options moving forward. Is TorchScript essentially deprecated? Will it break or be removed in the future or am I ok to keep relying on it? The only alternative to torchscript I see is to rewrite my code directly using the C++ API, but since I have users that rely on the python code, that would mean duplicating all of the codebase</p>

kilianyp: <p>AFAIK the new thing is AOTInductor:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://pytorch.org/tutorials/recipes/torch_export_aoti_python.html">
  <header class="source">

      <a href="https://pytorch.org/tutorials/recipes/torch_export_aoti_python.html" target="_blank" rel="noopener nofollow ugc">pytorch.org</a>
  </header>

  <article class="onebox-body">
    

<h3><a href="https://pytorch.org/tutorials/recipes/torch_export_aoti_python.html" target="_blank" rel="noopener nofollow ugc">torch.export AOTInductor Tutorial for Python runtime (Beta) — PyTorch...</a></h3>

  <p>An end-to-end example of how to use AOTInductor for Python runtime.</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>I am not 100% clear on the usage of autograd here, but I would expect to work in a similar fashion there.</p>
<p>There is <a href="https://pytorch.org/functorch/stable/notebooks/aot_autograd_optimizations.html" class="inline-onebox" rel="noopener nofollow ugc">AOT Autograd - How to use and optimize? — functorch nightly documentation</a></p>
<p>And as I understand it, that’s already being called by the <code>export</code> function so it might already work.</p>

IgnacioPickering: <p><code>torch.export</code> doesn’t seem to be able to export graphs with backward or <code>torch.autograd.grad</code> calls. A simple</p>
<pre><code class="lang-auto">import torch


class M(torch.nn.Module):
    def forward(self, x):
        x.sum().backward()
        return x.grad


torch.export.export(M(), (torch.tensor(1., requires_grad=True),))
</code></pre>
<p>fails withan unsupported operation error, so I don’t think autograd is a current use case of the AOTInductor. Compiled autograd seems to me to be exclusive to <code>torch.compile</code></p>
<p>I see what you mean however, it seems like <code>aot_function</code> can be used to get an fx-traced graph of the backward call. If this is something that torch.export can use to generate an exported program, then maybe AOTI can compile it into a blob? That would be a step in the right direction. If this is the case then I haven’t figure out how to pull that off yet. If the capability is there, it would be handy if <code>torch.export.export</code> had a <code>export_backward</code> flag =)</p>