2d sliding window implementation for numpy

sliding_window.py

import numpy as np
from numpy.lib.stride_tricks import as_strided
from typing import *

def sliding_window_2d(a: np.array, wnd_size: Tuple[int, int]):
    # a: an nd-array of (height, width, ...rest)
    # wnd_size: a tuple of (wnd_width, wnd_height)
    # returns an nd-array of (height - wnd_height + 1, width - wnd_width + 1, wnd_height, wnd_width, ...rest)
    height, width, *rest = a.shape
    s0, s1, *rest_strides = a.strides
    wnd_width, wnd_height = wnd_size
    return as_strided(a,
                      shape=(height - wnd_height + 1, width - wnd_width + 1, wnd_height, wnd_width, *rest),
                      strides=(s0, s1, s0, s1, *rest_strides))

if __name__ == '__main__':
    a = np.arange(0, 3 * 4).reshape((3, 4))
    print(a)
    print(sliding_window_2d(a, (2, 2)))
    assert (np.sum(sliding_window_2d(a, (2, 2)) == np.array([[2, 3], [6, 7]])) == 4)

    a = np.arange(0, 3 * 4 * 3).reshape((3, 4, 3))
    print(a)
    print(sliding_window_2d(a, (2, 2)))