Weighted sum of values
# grade this stage (from your course directory)
npx theinference test --stage weighted-sum --watchLast step: combine the value vectors using the attention weights. Four steps — score → mask → softmax → weighted sum — and you have a complete attention head. This single function is the literal engine of GPT-2: Chapter 6 runs it 12 times per block, 144 times across the whole model.
out[i] = Σ_j weights[i][j] · v[j]
q,k,v [T,d] ─▶ scores[T,T] ─▶ masked[T,T] ─▶ weights[T,T] ─▶ Σⱼ wᵢⱼ·v[j] ─▶ out[T,d]
What to build
Implement runAttention
Implement runAttention(q, k, v) (each [T, d]) returning the attention output [T, d]. You can compose the three functions you just wrote — runScores → runCausalMask → runSoftmaxRows — then take the weighted sum of v, or compute it directly.
export function runAttention(q: number[][], k: number[][], v: number[][]): number[][] {
// q, k, v: [T, d] -> [T, d]
}
The future positions are masked to ~0 weight, so they contribute nothing — you can skip them entirely (the reference does: for j in 0..=i). Either way is fine within tolerance.
Your output tensor ([T, d]) is compared to the reference with allclose, atol 1e-5. Inputs are randomized per run — you can't hardcode the answer.