Generating trace/compute graph

Hi, I’m wondering if there is any easy way to generate a high-level trace/compute graph for DGL scripts. I have tried using the torchviz/hiddenlayer libraries, but they are unable to include information from the graph update_all function (both the edge functions like copy, mul, etc and the aggregation function). Any ideas on how to get a good trace from DGL?


Basically all the operation on tensor are torch functions. There is custom op you can find at I’m not familiar with pytorch trace mechanism. One thing which worth noticing is that the custom op doesn’t only take tensor/strings, it also takes DGL’s data structure graph_index.