Reshaping a tensor of "block" sub-matrices from a stack to a row

There is probably a really simple way to do this but I’m just not seeing it, I want to obtain the following behavior using torch commands: I have a 2D tensor with dimension (Nn x n) , i.e. a tensor of N , nxn sub-matrices. I want to reshape it from being a stack of the block matrices (as it is now), into a horizontal row of these sub-matrices (dimension (n x Nn)) (i.e)

(A)
(B) —> (A B)

As an example if my tensor is (4, 2)
x=torch.tensor([[1,2],[3,4],[6,7],[8,9]])

tensor([[1, 2],
[3, 4],
[6, 7],
[8, 9]])

then the desired output would be a (2,4) of the following form

tensor([[1, 2, 6, 7],
[3, 4, 8, 9]])

where in this case the block sub-matrices are
A=[[1, 2], [3, 4]]
B=[[6, 7], [8, 9]]

Any help is always appreciated!

1 Like

If anyone has a similar problem, the solution is quite simple,

x=torch.tensor([[1,2],[3,4],[6,7],[8,9]])
print(torch.cat(torch.split(x,2),1))

tensor([[1, 2, 6, 7], [3, 4, 8, 9]])