B, S, E, H = 2, 5, 8, 2                  # batch, seq_len, embed_dim, heads

x = torch.randn(B, S, E)                 # torch.Size([2, 5, 8])
x = x.reshape(B, S, H, E // H)           # torch.Size([2, 5, 2, 4])  — split, don't move
x = x.permute(0, 2, 1, 3)                # torch.Size([2, 2, 5, 4])  — move heads forward