Attention Aggregation Lite

Problem

Perform one round of dot-product attention aggregation on an undirected graph.

For node v and neighbor w compute the raw score as the dot product of their feature vectors, apply a softmax over all neighbors of v, and return the weighted average of neighbor features. Isolated nodes return zeros. features is a list of equal-length vectors.

Hint

Subtracting the max score before exponentiating keeps the softmax stable.

Test cases

case 1

in: [3,[[0,1],[1,2]],[[1,0],[0,1],[1,1]]]

expected: [[0,1],[1,0.7310585786300049],[0,1]]

case 2

in: [4,[[0,1],[0,2],[0,3]],[[1,0],[1,0],[0,1],[0,-1]]]

expected: [[0.5761168847658291,0],[1,0],[1,0],[1,0]]

case 3

in: [2,[],[[1,2]]]

expected: [[0,0],[0,0]]

Starter codePython

import math
def attention_aggregation(n, edges, features):
    # Your code here
    pass

3 test cases · run your code against them in the editor