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 https://docs.dgl.ai/en/latest/tutorials/models/1_gnn/9_gat.html?highlight=visualization

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

Thank you for your responses.