Merging the heads back

After weighting V, each head’s output is (B, H, S, D_k). Transpose back before merging:

merged = head_values.transpose(0, 2, 1, 3).reshape(B, S, D)   # (2, 3, 6, 4) -> (2, 6, 12)

The test confirms this equals concatenating the heads along the feature axis: np.concatenate([head_values[:, h] for h in range(H)], axis=-1).

A trained layer would then apply an output projection W_O. The notebook stops at the concatenated heads.