✨[Feature] A hook function after each graph change or partition

[sean-xiang-applovin]

Is your feature request related to a problem? Please describe. I am trying to compile our model with tensorrt these days with exported program, our model is not very big, the original graph contains like 2k~ of nodes.

I am running into different problems, and there are problems even before we start to convert the partitioned graph, like during decomposition or post lowering phrase. Sometimes the debug log indicates some error of some node, and I cannot find that node in my original graph.

I find it helpful to dump the graph after each "change". Even though there is log printed, but

1. not all passes are printing the logs
2. it is hard to scroll back the log history and copy the whole graph of each pass

Describe the solution you'd like A hook function to be called after each intermediate graph, like graph after each pass, and each partitioned graph.

The function takes a gm object, and some metadata of this gm object, like a name, or a number k indicating it's the kth change to the original graph.

The most common implementation of this hook function is probably, to print the graph, or save it somewhere on the disk. So that I can check each graph and find which is the first graph contains the error node described in the error log.

I can also see this can be helpful if people want to visualize each graph and graph change during compilation.

Describe alternatives you've considered

Additional context