numpy rolling 2d

1def rolling_window(a, window_size):
2    shape = (a.shape[0] - window_size + 1, window_size) + a.shape[1:]
3    strides = (a.strides[0],) + a.strides
4    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)