Regarding the torch.export documentation, it is designed to trace the full graph at the Aten IR level, but potential failures in some corner cases. In such instances, is it possible to export multiple subgraphs, at the Aten IR level as well?
The motivation behind this is that the full graph might be too restrictive. Even if we only obtain the subgraph, there is still space for optimization for each of them.
Dynamo backends can be written as functions! One relatively straightforward approach here would be to pass a function to compile or torch._dynamo.optimize that then invokes aot_autograd (or any other torch ir-> aten ir step of your choosing) and records the functions in a list or structure of your choosing.
Note, that this isn’t really export anymore if you do it this way, since the very things that make export what it is seem to be different from what you want (whole graph) - but the step above would give you a single call way to extract all the graphs out. But it should give what you asked for.
Note, also, that in doing so, you silently are skipping whatever code couldn’t be captured and got delegated to eager. So, user beware!
Answered by mobile, apologies for being terse and any formatting issues.
Hi @voz, thanks a lot for getting back to me so quickly. I’ve had a quick look, and it is exactly what I was searching for. I’ve got another question regarding the transformation from torch IR to aten IR. Could you provide an example of how to perform these transformations or direct me to the main entry point in the code?
