Scaled Attention Scores

Problem

Compute scaled dot-product attention scores: Q @ K^T / sqrt(d_k) where d_k is the feature dimension of Q and K.

The signature is attention_scores(Q, K). Q has shape (seq_q, d_k), K has shape (seq_k, d_k) and the result has shape (seq_q, seq_k).

Hint

Dot each query row with every key row, then divide by the square root of the key dimension.

case 1

in: [[[1,0],[0,1]],[[1,0],[0,1]]]

expected: [[0.7071067812,0],[0,0.7071067812]]

case 2

in: [[[1,2]],[[1,2]]]

expected: [[3.5355339059]]

case 3

in: [[[1,0],[0,2]],[[1,0],[0,1],[1,1]]]

expected: [[0.7071067812,0,0.7071067812],[0,1.4142135624,1.4142135624]]

case 4

in: [[[2]],[[3]]]

expected: [[6]]

Starter codePython

import math
def attention_scores(Q, K):
    # Your code here
    pass

4 test cases · run your code against them in the editor