Sliding Window Attention Mask

Problem

Build a sliding window attention mask: position i attends to itself and the window - 1 positions before it. Entries the mask allows are 1 and masked entries are 0.

The signature is sliding_window_mask(seq_len, window). Return a seq_len x seq_len matrix.

Hint

A window of 1 degenerates to pure diagonal attention.

Test cases

case 1

in: [4,2]

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

case 2

in: [3,3]

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

case 3

in: [5,1]

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

case 4

in: [1,1]

expected: [[1]]

Starter codePython

def sliding_window_mask(seq_len, window):
    # Your code here
    pass

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