HardNLPnlp-364

Causal Masked Attention Row

Problem

In causal attention, position i may attend only to positions j <= i. Mask later positions with -1e9, then softmax.

Given the raw scores and the current row index, return the masked attention weights.

Hint

The masked logits shrink to zero weight after softmax.

case 1

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

expected: [1,0,0]

case 2

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

expected: [0.09003057317038046,0.24472847105479764,0.6652409557748218]

case 3

in: [[0,0],1]

expected: [0.5,0.5]

Starter codePython

def causal_masked_attention_row(scores, row_index):
    # Your code here
    pass

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