Visualizing attention heads in Graph Attention Network

Hi everyone,

I used the DGL to build a Graph Attention Network and I am play with adding some edge features. But I am interested in doing some of the visualizations in the original GAT paper and also in

Is there some tool or repo that I use to visualize various attention heads at various layers?

Thank you for your responses.