split_heads — reshape, then transpose
With H heads, split the D = H × D_k features into groups, then give the heads their own axis:
x.reshape(B, S, H, D // H).transpose(0, 2, 1, 3) # (B, S, D) -> (B, H, S, D_k)
| Stage | Shape | Meaning |
|---|---|---|
| Projected Q, K or V | (2, 6, 12) |
batch, token, features |
| Reshape | (2, 6, 3, 4) |
batch, token, head, head features |
| Transpose | (2, 3, 6, 4) |
batch, head, token, head features |
| Q @ Kᵀ | (2, 3, 6, 6) |
batch, head, query token, key token |
The permutation is (0, 2, 1, 3): axes 1 and 2 trade places, so the head comes before the token. The function rejects a head count that does not divide D, and booleans.