Export TorchScript Classes as TorchScript code

[alanhdu]

🚀 Feature

Add an option to export TorchScript classes as TorchScript code when torch.jit.save-ing a module. That way, the resulting file will be hermetically sealed, without relying on any particular code versioning to get it to execute.

Motivation

Right now, if you torch.jit.script a class, PyTorch relies on the actual class code to being available after torch.jit.load-ing the model for instance, in the code below I cannot actually execute my Module without having access to the A class

@torch.jit.script
class A:
    def __init__(self, x: int):
        self.x = x
class Module(torch.nn.Module):
    def forward(self):
        return A(1)
scripted = torch.jit.script(Module())
torch.jit.save(scripted, "test.pt")

This creates a really unfortunate coupling between the code defining the A class and the code need to run the Module class. This is also quite unintuitive, b/c scripted modules do not rely referencing code -- the code dependence seems only to beh ere for TorchScript non-module classes.

In our particular use case, we are using a TorchScript class as a "tuple but with some simple helper methods". Unfortunately, this dependence on our code is preventing us from refactoring our classes in our codebase (e.g. adding new fields) b/c of backwards compatibility with older models we've already serialized.

Pitch

Add an option to export the A class's code into the serialized torchscript module. This would allow torch.jit.script to produce hermetically sealed modules that don't need anything except PyTorch to execute correctly (which would massively simplify deployment for us).

Alternatives

If we don't do this, I think we should clearly document this behavior in the TorchScript classes section -- as I said earlier, b/c serializing a custom nn.Module does not require the underlying class definition, this code dependence came as a surprise to us.

cc @gmagogsfm

[tkuenzle]

This is the root cause for #42258. As @alanhdu mentioned, it works if the class is available when running the model. But this only works for Python and is a real blocker when running the model in another environment, eg. C++.

[byronyi]

Anyone’s working on this? If not I could take a stab.

[SplitInfinity]

I'm gradually polishing #44324 to close this.