Full self-attention from X

Background

Self-attention end to end: project X into Q, K, V, then run scaled dot-product attention.

$$Q,K,V = XW^Q, XW^K, XW^V, \qquad \text{out} = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

Problem statement

Implement self_attention(X, Wq, Wk, Wv) returning the attention output.

Input

• X — shape (n, d_model).
• Wq, Wk — (d_model, d_k); Wv — (d_model, d_v).

Output

Returns an np.ndarray of shape (n, d_v).

Examples

Example 1 — uniform attention averages the values

Input:  X = [[1, 0], [1, 0]], Wq = Wk = [[0, 0], [0, 0]], Wv = [[1, 0], [0, 1]]
Output: [[1, 0], [1, 0]]

Explanation: zero Q/K → all scores 0 → uniform weights → each output is the mean of the value rows ($V=[[1,0],[1,0]]$, mean $[1,0]$).

Constraints

• Project to Q, K, V; d_k = Q.shape[1].
• scores = Q @ K.T / sqrt(d_k); row-wise (numerically stable) softmax; multiply by V.
• Return shape (n, d_v).

Notes

• Every position is computed in one parallel matmul pipeline — no recurrence, and any token can influence any other in this single layer.