Project X into Q, K, V

Background

Self-attention starts by projecting the same input sequence X into three roles with learned matrices:

$$Q = XW^Q, \qquad K = XW^K, \qquad V = XW^V$$

The "self" is that all three come from one X — each token's vector is reused as a query, a key, and a value via different projections.

Problem statement

Implement qkv_projections(X, Wq, Wk, Wv) returning (Q, K, V).

Input

• X — shape (n, d_model): the input sequence.
• Wq, Wk, Wv — each shape (d_model, d_k).

Output

Returns (Q, K, V), each an np.ndarray of shape (n, d_k).

Examples

Example 1 — identity projections return X

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

Example 2 — different projections, different roles

Input:  X = [[1, 1]], Wq = [[1, 0], [0, 1]], Wk = [[2, 0], [0, 2]], Wv = [[0, 1], [1, 0]]
Output: Q = [[1, 1]], K = [[2, 2]], V = [[1, 1]]

Constraints

• Q = X @ Wq, K = X @ Wk, V = X @ Wv.
• Return the three matrices in that order.

Notes

• The same token row becomes three different vectors — which is exactly why Q/K/V are not the raw embeddings.