Edge-softmax with additional normalization dimension

Question regarding the edge-softmax function:
I am trying to implement this softmax operation


from the paper: https://arxiv.org/pdf/2109.05922.pdf.

Is there a way we can do this in dgl with/without using the edge-softmax function?

Hi, sorry for the late reply. It seems that the equation can be rewritten as:

\alpha^k_{viu} = \frac{\exp(\text{att}^k_{viu})}{\sum_{(v,j,z)\in \mathcal{E}}\exp(\text{att}^k_{vjz})}

Then this is exactly the edge-softmax operation. Let me know if you think my understanding is not correct.

You could use it with dgl.ops.edge_softmax and feed in a heterogeneous graph and a dictionary of edge logits. Although there are some bugs in the implementation ([Bug] A bunch of fixes in edge_softmax_hetero by BarclayII · Pull Request #4336 · dmlc/dgl · GitHub)

I agree, this operation becomes rather easy if we formulate the graph as heterogeneous with each relationship as a different edge. However, I wanted to achieve the same with a single edge (homogeneous) with the relationship features being encoded as a single feature vector per edge.

This is correct. But in my case v,j,z is not an independent edge. I have just one edge between each pair of linked nodes with the relational features being encoded as a feature vector for that edge. So, my edges are (v,x,z) where x is a vector of all relational features. N I need to normalize across all features for all neighbors. And edge softmax can normalize each feature dimension for an edge using all neighbors but not by all features for all neighbors. I guess heterograph is one of the ways wherein I can treat each relationship as an independent edge as u have suggested but is that the only way?