Is it possible to do a graph surgery on torchscript?

jit
1

Hello,
We are exploring the possibility of graph surgery on a jit scripted model.
The scripted model has a graph object. The graph object has blocks, nodes and variables.
If the graph has 10 nodes, is it possible to create two graph objects with 5 nodes each?

2 Likes
2

This is indeed possible, although if you are looking to do some lightweight graph transformations I would recommend torch.fx

3

Hello @Michael_Suo, thanks for your idea. Unfortunately we are bound to work on the JIT scripted model. So we can’t use fx for our purposes.

Do you think this is possible without fx using only the JIT API?

Thanks & Best regards,
RB

4

certainly it’s possible, you can look at some of the passes in torch/crsc/jit/passes as reference

5

Hi @Michael_Suo, I am thinking about how to auto add some fork/wait into the graph to enable the interop parallelism. Does the jit-trace/jit-script/FX already have this kind of feature or is it in the plan?