Swin Window Partition

Problem

Partition a 2D grid into non-overlapping window_size x window_size patches, scanning windows row-major and flattening each window row-major.

The signature is swin_window_partition(grid, window_size). Return the list of flattened windows.

Hint

Local windows cut attention cost from quadratic in image size to linear for a fixed window.

case 1

in: [[[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]],2]

expected: [[1,2,5,6],[3,4,7,8],[9,10,13,14],[11,12,15,16]]

case 2

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

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

case 3

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

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

case 4

in: [[[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]],4]

expected: [[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]]

Starter codePython

def swin_window_partition(grid, window_size):
    # Your code here
    pass

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