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?

Thanks

Basically all the operation on tensor are torch functions. There is custom op you can find at https://github.com/dmlc/dgl/blob/master/python/dgl/backend/pytorch/tensor.py#L288. 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.