Masked Attention Softmax

Problem

Apply a numerically stable softmax to each score row while honoring a boolean keep mask. Masked entries receive weight 0.0 and the surviving entries are renormalized to sum to 1.0.

The signature is masked_attention_softmax(scores, keep_mask). Every row must keep at least one entry.

Hint

Take the row maximum over kept entries only, then zero out the masked columns after exponentiating.

Test cases

case 1

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

expected: [[0.2689414214,0.7310585786],[0,1]]

case 2

in: [[[0,0,0]],[[1,0,1]]]

expected: [[0.5,0,0.5]]

case 3

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

expected: [[1]]

case 4

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

expected: [[0.5,0.5],[0.5,0.5]]

Starter codePython

import math
def masked_attention_softmax(scores, keep_mask):
    # Your code here
    pass

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