Sometimes that beautiful one-liner just isn't worth it compared to something like

torch.cat((
    img[:N, :N],
    img[N:, :N],
    img[:N, N:],
    img[N:, N:],
    ), dim=-1)

Skimage view_as_windows would do this.

This is spaceToDepth effectively? https://www.researchgate.net/figure/Intuitive-examples-of-space-to-depth-During-the-space-to-depth-method-top-we-aim-to_fig3_347156516

To accomplish the specific op here, spaceToDepth with block_size=2.

https://discuss.pytorch.org/t/is-there-any-layer-like-tensorflows-space-to-depth-function/3487/12

It can also be pixel unshuffle.Here you can find pytorch reference: https://pytorch.org/docs/stable/generated/torch.nn.PixelUnshuffle.html

I don't have the time to figure it out, but I'm pretty sure you can do it through some combination of permutations and reshapes. Play around with an NxN numpy array (e.g. np.arange(8**2).reshape(8,8)) and perform various transposes and reshapes and see what comes out. You might have to add and remove an axis at some point too.

Use convnet to transform to the right shape

Use convnet to transform to the right shape. May need to use dilution.