The diagram you know from the papers, the shape at every step, and the PyTorch that makes it — the same block three ways. Hover any box, shape, or line of code and the matching two light up.

(n, d)tokens in
→
(n, dₖ)×3Q K V
→
(n, n)scores
→
(n, n)weights A
→
(n, dₖ)output

the block (as the paper draws it)

Q K V MatMul QKᵀ Scale ÷√dₖ Mask (opt.) SoftMax MatMul ·V Output

the same block, in PyTorch

# x: your tokens                (n, d)
xQ, K, V = x@Wq, x@Wk, x@Wv     # (n, dₖ) eachscores = Q @ K.transpose(-2,-1) / dₖ**0.5   # (n, n)scores = scores.masked_fill(mask, -inf)     # (n, n)A = scores.softmax(dim=-1)     # (n, n) rows→1out = A @ V                    # (n, dₖ)

◆ derived shapes & dataflow from the standard attention definition  ·  ◇ illustrative n·dₖ are example sizes. Want a shape proven for every n? /research prove.