Repository: pytorch
Issue: FX should preserve type annotations and not break TorchScript-ability
Labels: oncall: jit,triaged,module: fx

## 🚀 Feature
Right now torch.fx tracing erases all python type annotations (on inputs or intermediate values), it breaks TorchScript that requires static typing.

## Motivation

One of the use cases of FX is to do program transformations from Python to Python that change the logic of the program only a little bit. Examples: quantization (local transformation) or model splitting (partitioning the model for pipelining or distributed inference according to some rules).

In those cases, the user might want to apply TorchScript to the transformed program (e.g. in order to run without Python, deploy to Mobile, achieve better perf, etc). FX should be as composible as possible with TorchScript for this use case and have good defaults.

## Pitch

I can list several ideas of the possible use cases we need to support (ranging from certain to more vague).

### No modifications roundtrip tracing

Claim: if a module X is both TorchScriptable (i.e. `torch.jit.script(X)` succeeds) and FX-able (i.e. `torch.fx.symbolic_trace(X)` succeeds and produces a python program with valid semantics) then `torch.jit.script(torch.fx.symbolic_trace(X))` should always succeed too.

In order to achieve this property I think it's sufficient to preserve type annotations for input arguments of `X.forward`. Semantics of FX are pretty restrictive data flow like, so most likely it's a subset of TorchScript type inference rules. However, it'd be good to try to prove it more formally. Specifically, TS sometimes needs to have annotations in the middle of the function (e.g. for an empty List of non-tensors). It might be that for valid FX programs it can never happen (I can't come up with any example of top of my head).

Possible implementation: just preserve `__annotations__` on the function and spit out them back in the codegen in GraphModule.

Simplest example
```
import torch
import torch.fx
from torch import Tensor
from typing import NamedTuple

class Pair(NamedTuple):
    x : Tensor
    y : Tensor

class Foo(torch.nn.Module):
    def forward(self, p : Pair, z : Tensor):
        return p.x+p.y+z

foo_scripted = torch.jit.script(Foo())
# works
foo_scripted(Pair(torch.rand(5), torch.rand(5)), torch.rand(5))

fxed = torch.fx.symbolic_trace(Foo())
# prints without type annotations :(
print(fxed.code)
# def forward(self, p, z):
#     getattr_1 = p.x
#     getattr_2 = p.y
#     add_1 = getattr_1 + getattr_2
#     add_2 = add_1 + z
#     return add_2

torch.jit.script(fxed)
# fails with:
# Tried to access nonexistent attribute or method 'x' of type 'Tensor (inferred)'.:
# def forward(self, p, z):
#     getattr_1 = p.x
#                 ~~~ <--- HERE
#     getattr_2 = p.y
#     add_1 = getattr_1 + getattr_2
```

### Simple re-splitting needs

Common FX passes might need to reorganize module hierarchy after flattening out. It's likely that intermediate values before the split were put on the module boundary, but were lost during flattening. It'd be nice to preserve them where possible by default.

In other words, passes akin to #44510 shouldn't break TorchScript-ability too.

### Advanced re-splitting needs

I'd argue that a basic `Node` class should have an optional `type_annotation` attribute (similarly to Python) which allows pass writers to introduce new nodes with non-Tensor types and still leverage standard GraphModule's codegen to produce a TorchScript-friendly program.

### Full-blown type inference

In general, if the module `X` is TorchScriptable, it has well-defined static types inferable for any intermediate values (because that's how TorchScript works). It'd be awesome if FX came with a type inference pass that can produce type annotations for intermediate values at least as powerful as TorchScript would. That'd ensure that pass writer can easily make their passes to produce TorchScriptable outputs as long as inputs are TorchScriptable.

From tech perspective, maintaining such pass might be hard though because of C++/python split. It also starts to get into MyPy-like territory.

cc @ezyang @SherlockNoMad @gmagogsfm