Mistake: the right shape means the right heads

“I can reshape directly to (B, H, S, D_k).”

example = np.arange(24).reshape(2, 3, 4)
correct = example.reshape(2, 3, 2, 2).transpose(0, 2, 1, 3)
wrong   = example.reshape(2, 2, 3, 2)
assert wrong.shape == correct.shape
assert not np.array_equal(wrong, correct)

Both arrays have the same shape. Only reshape-then-transpose keeps one token’s features inside one head. A direct reshape groups whatever sits next to it in memory, so head 0 of token 1 should be [4, 5] and the direct reshape hands it another token’s features.

Trace a value to expose the mismatch — a shape test cannot. Singleton axes can hide the mistake, because with only one token or one head the two orders coincide.