Attention step 2 — scale by √dₖ, then softmax

Background

The middle step turns raw scores into attention weights: divide by $\sqrt{d_k}$, then softmax along each query's row.

$$\alpha = \text{softmax}\!\left(\frac{S}{\sqrt{d_k}}\right) \quad\text{(row-wise)}$$

The $\sqrt{d_k}$ divisor keeps the logits from saturating as $d_k$ grows (large scores would push softmax to one-hot and kill gradients).

Problem statement

Implement attention_weights(scores, d_k) returning the row-wise softmax of the scaled scores.

Input

• scores — array-like of shape (n_q, n_k), raw $QK^\top$ scores.
• d_k — int, the key/query dimension.

Output

Returns an np.ndarray of shape (n_q, n_k) whose rows are non-negative and sum to 1.

Examples

Example 1 — the lesson's worked scores

Input:  scores = [[2, 0, 1]], d_k = 4
Output: [[0.50648039, 0.18632372, 0.30719589]]

Explanation: scaled by $\sqrt 4 = 2$ → $[1, 0, 0.5]$; softmax concentrates on the best match without zeroing the others.

Constraints

• Divide scores by sqrt(d_k) first.
• Softmax along the last axis (each query's row), numerically stable (subtract the row max).
• Each output row sums to 1.

Notes

• This is the only step where $d_k$ appears explicitly — and the reason transformers stay trainable as the head dimension grows.