Implement scaled dot-product attention

Task

In this interview you are asked to hand-write the forward pass of attention from the mathematical formula (no need to run code).

Implement single-head scaled dot-product attention.

Inputs

• Q: a 2D array of shape (Tq, d) (queries)
• K: a 2D array of shape (Tk, d) (keys)
• V: a 2D array of shape (Tk, dv) (values)
• Optional mask: a 2D array of shape (Tq, Tk) where mask[i][j] = 0/False means position j is not allowed to be attended to by query i.

Output

• O: a 2D array of shape (Tq, dv) defined by:

1. Compute scores:

$S = \frac{QK^\top}{\sqrt{d}} \quad \text{(shape } Tq \times Tk\text{)}$

2. Apply mask (if provided): for any masked-out entry S[i][j], treat it as $-\infty$ before softmax.

3. Apply row-wise softmax to get attention weights:

$A_{i,:} = \mathrm{softmax}(S_{i,:})$

4. Compute the output:

$O = AV$

Requirements / Discussion points

• Write clear pseudocode or Python-like code for the above.
• Explain how you would implement numerically stable softmax.
• State the time complexity in terms of Tq, Tk, d, and dv.
• Mention at least two edge cases (e.g., fully-masked row, large magnitudes, shape mismatches).

Implement Scaled Dot-Product Attention

Implement the forward pass of single-head scaled dot-product attention.

Inputs

• Q: a 2D list of shape (Tq, d) (queries)
• K: a 2D list of shape (Tk, d) (keys)
• V: a 2D list of shape (Tk, dv) (values)
• mask (optional): a 2D list of shape (Tq, Tk) where a falsy entry mask[i][j] (0/False) means key position j is not allowed to be attended to by query i. If mask is None, all positions are allowed.

Output

Return O, a 2D list of shape (Tq, dv):

1. Scores: S = (Q · Kᵀ) / sqrt(d), shape (Tq, Tk).
2. Mask: for any masked-out entry, set S[i][j] = -inf before softmax.
3. Softmax (row-wise): A[i,:] = softmax(S[i,:]) using the numerically stable form (subtract the row max before exponentiating).
4. Output: O = A · V.

Conventions for this challenge

• Use a numerically stable softmax: subtract the maximum finite score in the row before calling exp. Masked entries (-inf) contribute weight 0.
• For a fully-masked row (every key disallowed), output a uniform distribution over all Tk keys (weight 1/Tk each) rather than producing NaN.
• If Q is empty (Tq == 0), return [].
• Round every output value to 6 decimal places so results are deterministic.

Constraints

• 1 <= d, dv (each row of Q/K has length d; each row of V has length dv)
• 0 <= Tq, Tk (Tq == 0 returns an empty list)
• K and V share the same first dimension Tk
• mask, when provided, has shape (Tq, Tk); a falsy entry forbids attending to that key
• All numeric inputs are real-valued floats; output values are rounded to 6 decimal places

Examples

Example 1

Input: ([[1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]], [[1.0, 2.0], [3.0, 4.0]], None)
Output: [[1.660477, 2.660477]]

Explanation: Single query [1,0], d=2 so scale=sqrt(2). Raw scores = [1,0]; scaled = [0.7071, 0]. Softmax ~ [0.6698, 0.3302]; output = 0.6698*[1,2] + 0.3302*[3,4] = [1.660477, 2.660477].

Example 2

Input: ([[0.0, 0.0]], [[5.0, -3.0], [2.0, 8.0]], [[10.0, 0.0], [0.0, 10.0]], None)
Output: [[5.0, 5.0]]

Explanation: Query is all zeros, so every score is 0 regardless of keys; softmax gives uniform weights [0.5, 0.5]; output = 0.5*[10,0] + 0.5*[0,10] = [5,5].

Hints

1. Process one query row at a time: compute its Tk scores, softmax them, then combine the value rows.
2. Numerically stable softmax: subtract the row's maximum finite score before exponentiating; this prevents overflow for large magnitudes and leaves the result unchanged.
3. Treat a masked entry as score -inf, which exponentiates to weight 0 — but guard the fully-masked row separately (all -inf) to avoid dividing by zero; here we fall back to a uniform distribution.
4. Don't forget the 1/sqrt(d) scaling on the raw dot products before softmax.