[FIXED] Quick method to manually expand a PyTorch tensor

Issue

I need a quick and simple method to transform a PyTorch tensor with dimensions (D, M, M) into a tensor with dimensions (D // 4, M x 2, M x 2) manually without convolutions. I previously arrived to a solution for the oposite operation.(see comments on the previous question (oposite operation) )

previous question.

Quick method to transform a PyTorch tensor with gradient preservation

# Batch, Depth ,size, size -> Batch, 4 x Depth ,size//2, size//2
x = x.unfold(2, 2, 2).unfold(3, 2, 2).reshape(batch, depth, size // 2, size // 2, 4)
x = x.permute(0, 2, 3, 1, 4).reshape(batch, size // 2, size // 2, depth * 4).permute(0, 3, 1, 2)

Desired Input:

Transform (12, 2, 2) to (3, 4, 4):

[[[ 0,  1,  4,  5, 16, 17, 20, 21, 32, 33, 36, 37],

[ 2,  3,  6,  7, 18, 19, 22, 23, 34, 35, 38, 39]],

[[ 8,  9, 12, 13, 24, 25, 28, 29, 40, 41, 44, 45],

[10, 11, 14, 15, 26, 27, 30, 31, 42, 43, 46, 47]]]

Desired Output:

[[[ 0,  1,  2,  3], [ 4,  5,  6,  7], [ 8,  9, 10, 11], [12, 13, 14, 15]],

[[16, 17, 18, 19], [20, 21, 22, 23], [24, 25, 26, 27], [28, 29, 30, 31]],

[[32, 33, 34, 35], [36, 37, 38, 39], [40, 41, 42, 43], [44, 45, 46, 47]]]

EDIT: preferable working with batch. test code:

import torch
batch = 2
input = torch.arange(3*4*4*batch).reshape(batch,3, 4, 4)
batch, depth, size, _ = input.shape
x = input.unfold(2, 2, 2).unfold(3, 2, 2).reshape(batch, depth, size // 2, size // 2, 4)
x = x.permute(0, 2, 3, 1, 4).reshape(batch, size // 2, size // 2, depth * 4).permute(0, 3, 1, 2)
**reverse operation code**
print((input==x).all())

Solution

SOLUTION:

x.permute(0, 2, 3, 1).reshape(batch, size//2, size//2, depth, 
2,2).permute(0, 3, 1, 4, 2, 5).reshape(batch, depth, size, size)

Answered By - MasterEpilif