Reshape x from (B, T, D) into (B, H, T, D/H) where H is num_heads and D divides evenly.

Split the feature dimension first and then move the head axis in front of time, so each head ends up with a contiguous slice of the features. Reshaping straight to the target shape would interleave features across heads instead.

Input

x =
tensor([[[ 0.,  1.,  2.,  ...,  5.,  6.,  7.],
         [ 8.,  9., 10.,  ..., 13., 14., 15.],
         [16., 17., 18.,  ..., 21., 22., 23.]]])
num_heads = 2

Output

tensor([[[[ 0.,  1.,  2.,  3.],
          [ 8.,  9., 10., 11.],
          [16., 17., 18., 19.]],

         [[ 4.,  5.,  6.,  7.],
          [12., 13., 14., 15.],
          [20., 21., 22., 23.]]]])