Okay, now that I know more about PT I can answer my own question.
The solution is to use TorchScript and work directly with the traced graph.
For that you`ll need a torch::jit backend.
Here`s a useful article 5 on the topic. Bear in mind however that it`s already quite outdated, so e.g. CustomFuseGraph would not be of much help so you`ll need to create fusion groups (i.e. meta-primitives you`ll be in charge of compiling and executing) by hand — substituting the existing prims via SubgraphUtils::createSingletonSubgraphAndUpdateAliasing() and SubgraphUtils::mergeNodeIntoSubgraphAndUpdateAliasing().