Attention over edges in Graph Neural Network

My question is kind of puzzle. The question is related to Graph Neural Network. What I want to do is attention over adjoint to the same node edges. So what I got is tensor of edges features (size = number_of_edges x number_of_fetures) and tensor of node indecies (size = number_of_edges) indicating to what node particular adge is adjointed. For example:

edges_features = torch.randint(3, (5, 3))
edges_features
tensor([[1, 2, 1],
        [2, 2, 1],
        [2, 2, 1],
        [0, 1, 2],
        [2, 0, 2]])

here we have 5 edges with 3 features per each

node_index = torch.randint(2, (5,)) 
node_index
tensor([0, 1, 0, 0, 1])

two nodes. 0th node has 3 adjoin edge, 1st have 2 adjoin edge

For simplification, I consider key, query, and value as edge features. So what I need to do is calculate pairwise dot products of edge features correspond to the same node (for the 0th node it will be tensor 3x3, for the 1st matrix 2x2), then make softmax over results (row-wise) and after calculating a weighted sum of edge features edgewise. As result, I expect a tensor of the same size as the edges_features tensor.
Any suggestions?