Length-Aware Causal Mask

Problem

Apply a causal mask that also respects per-row valid lengths: entry (i, j) is kept only when j <= i and j < lengths[i], otherwise it becomes -1e9.

The signature is causal_mask_with_lengths(scores, lengths). This is used for packed or padded batches where rows hold different numbers of real tokens.

Hint

Both conditions matter: the row length can mask out positions that the causal rule would allow.

Test cases

case 1

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

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

case 2

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

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

case 3

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

expected: [[5]]

case 4

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

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

Starter codePython

def causal_mask_with_lengths(scores, lengths):
    # Your code here
    pass

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