Multi-head attention shape flow

Makes the attention equation type-checkable by exposing every tensor axis and the contracted dimensions.

Q K and V have batch, head, token, and head-width axes; attention weights have batch, head, query-token, and key-token axes; multiplying by V restores a final head-width axis.

Within every batch and head, compare every query token with every key token, then use those weights to combine value features.

Tensor contraction is the structural verb; each matrix product removes one named shared axis and preserves the others.

Real implementations may use different query and key lengths, grouped KV heads, packed sequences, or fused layouts while preserving the same contraction logic.