What is vectorised way of doing this operation in pytorch instead of two FOR loops


Hello,
I have a tensor 'A' in Pytorch of dimesnsions Batch x Channel x Height x Width. I want to reshape it into 'B' such that dimesnions H and W are increased by 'r' and channels reduced by a factor of 'r^2'. For 'r'=2, the illustration is shown in figure attached.
In the figure if 'B' had 4 channels then first 4 channels, with 4th channel having border color violet, then first four channels in 'A' will be peack/skin colored pixels with border colors red, green, blue and violet and remaining channels adjusted accordingly.

I know the 'pack' and 'unpack' each can be done with 2 for loops. But that takes more time. There shoud be vectorised way in PyTorch to do switch between 'A' and 'B' with just reshape and pemutation commands. Can someone help me on that?

In this example, the batch size is set to 1. But if batch dimension is more I would like the operations shown in the figure to operate individually on each batch entry.

Can someone please help me with a generic code to switch between A and B in the vectorised way in PyTorch. Better if that also works when Batch size is more than 1.

Please note the two operations cannot be done with already implemented nn.PixelShuffle

Thanks a lot.  

Although this can be done with careful permutation and reshaping, pytorch has already implemented this with nn.PixelShuffle.
