Functionalization in PyTorch: Everything You Wanted To Know

compiler
3

Yes, this feature is available in the 2.0 release. torch.compile() uses functionalization internally, and it’s also available more directly as a program transform using functorch.functionalize().