A single attention head·8 min

Softmax the rows

0 / 1

# grade this stage (from your course directory)
npx theinference test --stage softmax-rows --watch

Turn each row of (masked) scores into attention weights — non-negative numbers that sum to 1.

Softmax, applied per row

weights[i] = softmax(scores[i])
  where softmax(s)[j] = exp(s[j] - max(s)) / Σ_k exp(s[k] - max(s))

Worked example: a masked row

Take the row [2, 0, -1e9, -1e9] (two real scores, two future positions masked):

Subtract the row max (2) → [0, -2, -1e9, -1e9]
exp → [1, 0.1353, ~0, ~0]
sum = 1.1353
normalize → [0.8808, 0.1192, 0, 0]

The two -1e9 entries vanished to 0 automatically — that's the mask doing its job through softmax. (Same numbers as position 1 in the very first worked example.)

What to build

Your task

Implement runSoftmaxRows

Implement runSoftmaxRows(rows): apply softmax to each row independently. Subtract each row's max before exp — mathematically identical, but it prevents overflow.

tests/adapters.ts

export function runSoftmaxRows(rows: number[][]): number[][] {
  // [r, c] -> [r, c]; each row sums to 1
}

Subtract the row max FIRST

Without it, exp(large) overflows to Infinity and you get Infinity / Infinity = NaN. Subtracting the max is mathematically identical (the shift cancels in the ratio) and is the standard numerically-stable softmax.

How it's graded

Your output (same shape as the input) is compared with allclose, atol 1e-5. Each row should sum to 1.

Why mask the future Weighted sum of values