How to convert this 3d matrix to 2d efficiently?

I have a matrix A which is 3d and I want to convert so that it is equal to B which is a 2d matrix

> A = torch.tensor( [ [[1,1,1,1,1],
>          [2,2,2,2,2],
>          [3,3,3,3,3],
>          [4,4,4,4,4],
>          [ 5,5,5,5,5]],
> 
>         [[6,6,6,6,6]
>          [7,7,7,7,7],
>          [8,8,8,8,8],
>          [9,9,9,9,9],
>          [ 10,10,10,10,10]],
> 
>         [[11,11,11,11,11],
>          [12,12,12,12,12],
>          [13,13,13,13,13],
>          [14,14,14,14,14],
>          [ 15,15,15,15,15]]])
> 
> B = torch.tensor( [[1,1,1,1,1, 6,6,6,6,6, 11,11,11,11,11],
>          [2,2,2,2,2, 7,7,7,7,7, 12,12,12,12,12],
>          [3,3,3,3,3, 8,8,8,8,8, 13,13,13,13,13],
>          [4,4,4,4,4, 9,9,9,9,9, 14,14,14,14,14],
>          [ 5,5,5,5,5, 10,10,10,10,10, 15,15,15,15,15]])

Please help

I’m really awful at working out dimensions from printing out arrays. But if I’m right A.shape = (5,5,3) and B.shape = (5,15)? If this is the case could you not just do A = torch.reshape(A,B.shape())

1 Like

Directly reshaping A will unfortunately not work, as it would result in:

print(A.view(B.size()))
> tensor([[ 1,  1,  1,  1,  1,  2,  2,  2,  2,  2,  3,  3,  3,  3,  3],
          [ 4,  4,  4,  4,  4,  5,  5,  5,  5,  5,  6,  6,  6,  6,  6],
          [ 7,  7,  7,  7,  7,  8,  8,  8,  8,  8,  9,  9,  9,  9,  9],
          [10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 12, 12, 12, 12, 12],
          [13, 13, 13, 13, 13, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15]])

To get the desired output you would need to permute the first two dimensions first:

print(A.permute(1, 0, 2).reshape(B.size()))
> tensor([[ 1,  1,  1,  1,  1,  6,  6,  6,  6,  6, 11, 11, 11, 11, 11],
          [ 2,  2,  2,  2,  2,  7,  7,  7,  7,  7, 12, 12, 12, 12, 12],
          [ 3,  3,  3,  3,  3,  8,  8,  8,  8,  8, 13, 13, 13, 13, 13],
          [ 4,  4,  4,  4,  4,  9,  9,  9,  9,  9, 14, 14, 14, 14, 14],
          [ 5,  5,  5,  5,  5, 10, 10, 10, 10, 10, 15, 15, 15, 15, 15]])
3 Likes