Queries against keys

Attention starts with one line, and the line is entirely about shapes.

q  (batch, q_len, d)
k  (batch, k_len, d)

want: scores (batch, q_len, k_len)
      scores[b, i, j] = dot(q[b, i], k[b, j]) / sqrt(d)

Every query against every key. The d axis is summed away, and two length axes that were both position axes become rows and columns of a matrix.

Getting there

@ contracts the last axis of the left with the second to last of the right. q has d last; k also has d last. So k has to be transposed:

q @ k.transpose(-2, -1)      # (b, q_len, d) @ (b, d, k_len) -> (b, q_len, k_len)

transpose(-2, -1) rather than transpose(1, 2), so the line survives a head axis being added in front. It is a view, so it costs nothing.

The division

scores / (d ** 0.5)

Each score is a sum of d products. If the entries of q and k are roughly unit-scale and independent, that sum has variance proportional to d, so its standard deviation grows like sqrt(d). Without the division the scores spread wider as the model gets wider, the softmax saturates, and the gradients vanish before training starts.

Dividing by sqrt(d) holds the spread constant regardless of head width. That is the whole reason it is there, and it is why it is sqrt(d) rather than d.

Your task

def scores(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor

Return (batch, q_len, k_len) scaled dot-product scores. No softmax, no mask; just the scores.

The starter forgets to transpose and multiplies the wrong axes together, which raises unless q_len, k_len and d happen to agree.