Easy

Attention Context Vector

Easy

~10 min

code completion

After converting raw scores to attention weights $α$ (via softmax), the final attention output is a weighted sum of the value vectors:

C = α V

α

has shape

(T_{q}, T_{k})

— attention weights; each row sums to 1

V

has shape

(T_{k}, d_{v})

— value matrix

C

has shape

(T_{q}, d_{v})

— context vectors

Each row $C_{i}$ is a convex combination of the rows of $V$ , weighted by how strongly query $i$ attended to each key.

For a single query, $α$ may be a 1D vector of length $T_{k}$ , and the output is a 1D context vector of length $d_{v}$ .

Your task:

Implement attention_context(weights, V) that returns $α @ V$ .

Example Tests

Single query attending fully to first value row

Input: {"V":[[1,2],[3,4],[5,6]],"weights":[1,0,0]}

Expected: [1,2]

Equal weights over two value rows: midpoint

Input: {"V":[[2,4],[6,8]],"weights":[0.5,0.5]}

Expected: [4,6]

Uniform weights: output is mean of value rows

Input: {"V":[[3,6],[6,3],[9,3]],"weights":[0.3333333333333333,0.3333333333333333,0.3333333333333333]}

Expected: [6,4]

You can read the full problem statement above. Create a free account to run code in the browser, submit solutions, and track your progress.