Undo the head split: turn x of shape (B, H, T, Dh) back into (B, T, H*Dh), so each time step's heads are concatenated in head order.

Moving the axes before reshaping is what makes this the exact inverse of the split; reshaping first would scramble which head each feature came from.

Input

tensor([[[[ 0.,  1.,  2.,  3.],
          [ 4.,  5.,  6.,  7.],
          [ 8.,  9., 10., 11.]],

         [[12., 13., 14., 15.],
          [16., 17., 18., 19.],
          [20., 21., 22., 23.]]]])

Output

tensor([[[ 0.,  1.,  2.,  ..., 13., 14., 15.],
         [ 4.,  5.,  6.,  ..., 17., 18., 19.],
         [ 8.,  9., 10.,  ..., 21., 22., 23.]]])