Window Attention Shift Mask

Problem

Build the attention mask for shifted-window (Swin) attention over one window.

Tokens are laid out row-major on a window_size x window_size grid. For tokens a and b output 0 when both lie on the same side of the shift in rows and columns, otherwise 1. Return the (window_size^2) x (window_size^2) mask of 0/1 values.

Test cases

case 1

in: [2,1]

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

case 2

in: [2,0]

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

case 3

in: [1,0]

expected: [[0]]

Starter codePython

def shift_window_mask(window_size, shift):
    # Your code here
    pass

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