Funtions similiar to `edge_softmax`

there exists a line using function called edge_softmax in gat pytorch example ( https://github.com/dmlc/dgl/blob/828a5e5bc6ffaa5716b02283874ec830e1b786fc/python/dgl/nn/pytorch/conv/gatconv.py )

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.