I am interested in dividing a tensor to subtensors.
eg. if i have tensor A = torch.rand(2,3,8,8) i what to generate tenosr B which has the size of 2x3x16x2x2 and then convert it back to in a way that has the same size as A.
in  this case, the dimension of tensor A is divided by k = 4 and we chose subblock size of s=8//k=2 (2x2).
in other words, the subblock are chosen in the following way when the same color means those elements are belonging to the same subblock

image.png625×690 1.88 KB

im not sure how to do it in a fast way, (the brutforce way is using looops), i think the fold unfold should be faster.
can anyone help me please how i can code it in a fast way?