torch.export code that includes torch.autograd.grad

[Linux-cpp-lisp]

🐛 Describe the bug

The following minimal example (based on a large real-world model which fails the same way) fails with errors in torch.export:

import torch

class F(torch.nn.Module):
    def forward(self, x):
        y = x.square()
        return torch.autograd.grad(y.sum(), x)

f = F()
x = torch.ones(3)
x.requires_grad_(True)
print(f(x))

f = F()
x = torch.ones(3)
x.requires_grad_(True)
print(torch.export.export(f, args=(x,), strict=False))

printing (omitting backtraces):

(tensor([2., 2., 2.]),)
[...]
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

With strict=True, it errors as

Unsupported: 'skip function grad in file site-packages/torch/autograd/__init__.py'

Is exporting through a module that includes internal calls to the autograd engine a supported usecase for torch.export (and other related APIs, like torch._export.aot_compile)?

Error logs

No response

Minified repro

No response

Versions

PyTorch 2.3.0

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

[bdhirsh]

from @avikchaudhuri : seems reasonable to have a better error message for export when people use autograd.grad or .backward() directly in the code they are exporting

[Linux-cpp-lisp]

@bdhirsh does that mean that code containing auto-differentiation is not supported for torch.export? (Is this in the roadmap, or not intended to work?)

[tugsbayasgalan]

cc: @JackCaoG for the dynamo tracing through autograd.