rolling window 2d array

1# Rolling window for 2D arrays in NumPy
2import numpy as np
3
4def rolling_window(a, shape):  # rolling window for 2D array
5    s = (a.shape[0] - shape[0] + 1,) + (a.shape[1] - shape[1] + 1,) + shape
6    strides = a.strides + a.strides
7    return np.lib.stride_tricks.as_strided(a, shape=s, strides=strides)