Causal Mask

Problem

Apply a causal (lower-triangular) mask to an attention score matrix: keep entries on or below the diagonal and replace every entry strictly above the diagonal with -1e9.

The signature is causal_mask(scores). Scores may be rectangular; row i masks columns j > i.

Hint

Compare the column index with the row index; -1e9 acts as a masked-out score.

Test cases

case 1

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

expected: [[1,-1000000000],[3,4]]

case 2

in: [[[0]]]

expected: [[0]]

case 3

in: [[[1,2,3],[4,5,6],[7,8,9]]]

expected: [[1,-1000000000,-1000000000],[4,5,-1000000000],[7,8,9]]

case 4

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

expected: [[1,-1000000000,-1000000000],[4,5,-1000000000]]

Starter codePython

def causal_mask(scores):
    # Your code here
    pass

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