Attention Weighted Sum

Problem

Compute the attention output as the weighted sum of value vectors: output[i] = sum over j of weights[i][j] * V[j].

The signature is attention_weighted_sum(weights, V) where weights has shape (seq_q, seq_k) and V has shape (seq_k, d_v).

Hint

Each output row is a convex combination of the value rows weighted by the attention row.

case 1

in: [[[0.5,0.5]],[[1,2],[3,4]]]

expected: [[2,3]]

case 2

in: [[[1,0,0]],[[1,2],[3,4],[5,6]]]

expected: [[1,2]]

case 3

in: [[[0.25,0.25,0.5]],[[1,0],[0,1],[1,1]]]

expected: [[0.75,0.75]]

case 4

in: [[[1],[0]],[[5,6]]]

expected: [[5,6],[0,0]]

Starter codePython

def attention_weighted_sum(weights, V):
    # Your code here
    pass

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