Attention Weight Sum

Problem

Sum the k largest attention weights, the top-k attention mass. Sort the weights in descending order and sum the first k; when k exceeds the number of weights, use all of them.

Return 0.0 for empty weights or k <= 0.

Hint

A high top-k mass means attention concentrates on few timesteps.

Test cases

case 1

in: [[0.5,0.3,0.2],1]

expected: 0.5

case 2

in: [[0.4,0.4,0.2],2]

expected: 0.8

case 3

in: [[1],2]

expected: 1

case 4

in: [[],3]

expected: 0

Starter codePython

def attention_weight_sum(weights, k):
    # Your code here
    pass

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