Top-K Gradient Compression

Problem

Compress a gradient matrix to its top-k entries by absolute value. Return the selected entries as [row, col, value] sorted by descending absolute value, breaking ties by smaller row and then smaller column.

The signature is topk_compress(gradients, k). Return at most k entries.

Hint

Top-k sparsification is the simplest gradient compression used in distributed training.

Test cases

case 1

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

expected: [[1,1,4],[1,0,3]]

case 2

in: [[[1,-5],[2,0.5]],1]

expected: [[0,1,-5]]

case 3

in: [[[0,0],[0,0]],2]

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

case 4

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

expected: [[0,1,2],[0,0,1]]

case 5

in: [[[1]],0]

expected: []

Starter codePython

def topk_compress(gradients, k):
    # Your code here
    pass

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