Attention step 1 — the scores QKᵀ

Background

The first step of attention is scoring: each query is dot-producted with every key to measure content similarity.

$$S = QK^\top, \qquad S_{ij} = \mathbf{q}_i \cdot \mathbf{k}_j$$

Q is (n_q, d_k), K is (n_k, d_k), and S is (n_q, n_k) — one raw similarity per query-key pair.

Problem statement

Implement attention_scores(Q, K) returning the raw score matrix $QK^\top$.

Input

• Q — array-like of shape (n_q, d_k).
• K — array-like of shape (n_k, d_k).

Output

Returns an np.ndarray of shape (n_q, n_k).

Examples

Example 1 — the lesson's worked query

Input:  Q = [[1, 0, 1, 0]],
        K = [[1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 0, 0]]
Output: [[2, 0, 1]]

Explanation: $q\cdot k_1=2$, $q\cdot k_2=0$, $q\cdot k_3=1$.

Constraints

• Compute Q @ K.T.
• Return shape (n_q, n_k).

Notes

• A larger score means the query and key point in more similar directions — the content match attention is about to turn into a weight.