Funtions similiar to `edge_softmax`

there exists a line using function called edge_softmax in gat pytorch example ( )

graph.edata['a'] = self.attn_drop(edge_softmax(graph, e))

however, I want to know, is there any function similiar to this, such as a simple normalization, edge_norm. I have tried to imitate edge_softmax source code, however, it is wrapped too highly to read.

I think the latest version of edge_softmax is much easier to read:

which simply uses the spmm and sddmm kernels we support recently.