[FIXED] Reshape tensors in pytorch?

Issue

I'm struggling with the result of a matrix multiplication in pytorch and I don't know how to solve it, in particular: I'm multiplying these two matrices

tensor([[[[209.5000, 222.7500],
          [276.5000, 289.7500]],

         [[208.5000, 221.7500],
          [275.5000, 288.7500]]]], dtype=torch.float64)

and

tensor([[[[ 0.,  1.,  2.,  5.,  6.,  7., 10., 11., 12.],
          [ 2.,  3.,  4.,  7.,  8.,  9., 12., 13., 14.],
          [10., 11., 12., 15., 16., 17., 20., 21., 22.],
          [12., 13., 14., 17., 18., 19., 22., 23., 24.]],

         [[25., 26., 27., 30., 31., 32., 35., 36., 37.],
          [27., 28., 29., 32., 33., 34., 37., 38., 39.],
          [35., 36., 37., 40., 41., 42., 45., 46., 47.],
          [37., 38., 39., 42., 43., 44., 47., 48., 49.]],

         [[50., 51., 52., 55., 56., 57., 60., 61., 62.],
          [52., 53., 54., 57., 58., 59., 62., 63., 64.],
          [60., 61., 62., 65., 66., 67., 70., 71., 72.],
          [62., 63., 64., 67., 68., 69., 72., 73., 74.]]]],
       dtype=torch.float64)

with the following line of code A.view(2,-1) @ B, and then I reshape the result with result.view(2, 3, 3, 3). The resulting matrix is

tensor([[[[ 6687.5000,  7686.0000,  8684.5000],
          [11680.0000, 12678.5000, 13677.0000],
          [16672.5000, 17671.0000, 18669.5000]],

         [[ 6663.5000,  7658.0000,  8652.5000],
          [11636.0000, 12630.5000, 13625.0000],
          [16608.5000, 17603.0000, 18597.5000]],

         [[31650.0000, 32648.5000, 33647.0000],
          [36642.5000, 37641.0000, 38639.5000],
          [41635.0000, 42633.5000, 43632.0000]]],


        [[[31526.0000, 32520.5000, 33515.0000],
          [36498.5000, 37493.0000, 38487.5000],
          [41471.0000, 42465.5000, 43460.0000]],

         [[56612.5000, 57611.0000, 58609.5000],
          [61605.0000, 62603.5000, 63602.0000],
          [66597.5000, 67596.0000, 68594.5000]],

         [[56388.5000, 57383.0000, 58377.5000],
          [61361.0000, 62355.5000, 63350.0000],
          [66333.5000, 67328.0000, 68322.5000]]]], dtype=torch.float64)

Instead I want

tensor([[[[ 6687.5000,  7686.0000,  8684.5000],
          [11680.0000, 12678.5000, 13677.0000],
          [16672.5000, 17671.0000, 18669.5000]],

         [[31650.0000, 32648.5000, 33647.0000],
          [36642.5000, 37641.0000, 38639.5000],
          [41635.0000, 42633.5000, 43632.0000]],

         [[56612.5000, 57611.0000, 58609.5000],
          [61605.0000, 62603.5000, 63602.0000],
          [66597.5000, 67596.0000, 68594.5000]]],


        [[[ 6663.5000,  7658.0000,  8652.5000],
          [11636.0000, 12630.5000, 13625.0000],
          [16608.5000, 17603.0000, 18597.5000]],

         [[31526.0000, 32520.5000, 33515.0000],
          [36498.5000, 37493.0000, 38487.5000],
          [41471.0000, 42465.5000, 43460.0000]],

         [[56388.5000, 57383.0000, 58377.5000],
          [61361.0000, 62355.5000, 63350.0000],
          [66333.5000, 67328.0000, 68322.5000]]]], dtype=torch.float64)

Can someone help me? Thanks

Solution

This is a common but interesting problem because it involves a combination of torch.reshapes and torch.transpose to solve it. More specifically, you will need

1. Apply an initial reshape to restructure the tensor and expose the axes you want to swap;
2. Then do so using a transpose operation;
3. Lastly apply a second reshape to get to the desired format.

In your case, you could do:

>>> result.reshape(3,2,3,3).transpose(0,1).reshape(2,3,3,3)
tensor([[[[ 6687.5000,  7686.0000,  8684.5000],
          [11680.0000, 12678.5000, 13677.0000],
          [16672.5000, 17671.0000, 18669.5000]],

         [[31650.0000, 32648.5000, 33647.0000],
          [36642.5000, 37641.0000, 38639.5000],
          [41635.0000, 42633.5000, 43632.0000]],

         [[56612.5000, 57611.0000, 58609.5000],
          [61605.0000, 62603.5000, 63602.0000],
          [66597.5000, 67596.0000, 68594.5000]]],


        [[[ 6663.5000,  7658.0000,  8652.5000],
          [11636.0000, 12630.5000, 13625.0000],
          [16608.5000, 17603.0000, 18597.5000]],

         [[31526.0000, 32520.5000, 33515.0000],
          [36498.5000, 37493.0000, 38487.5000],
          [41471.0000, 42465.5000, 43460.0000]],

         [[56388.5000, 57383.0000, 58377.5000],
          [61361.0000, 62355.5000, 63350.0000],
          [66333.5000, 67328.0000, 68322.5000]]]], dtype=torch.float64)

I encourage you to look a the intermediate results to get an idea of how the method works so you can apply it on other use cases in the future.

Answered By - Ivan