Attention Head Pruning Top-K

Problem

Select the top-k attention heads to keep when pruning, using their importance scores; break ties by preferring the smaller head index.

The signature is head_prune_topk(importance, k). Return the kept head indices sorted ascending.

Hint

Sort by descending score and ascending index, keep the first k, then sort the result for a stable answer.

case 1

in: [[0.5,0.9,0.1,0.9],2]

expected: [1,3]

case 2

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

expected: [0,1,2]

case 3

in: [[0.2,0.2,0.2],1]

expected: [0]

case 4

in: [[5,4,3,2],2]

expected: [0,1]

Starter codePython

def head_prune_topk(importance, k):
    # Your code here
    pass

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