Trajectory Transformer Attention

Problem

Compute softmax attention weights for one query over a list of key vectors.

Scores are dot(query, key) / temperature and the softmax uses the max-subtraction trick. Keys is non-empty.

Hint

Low temperature concentrates attention on the most similar key.

case 1

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

expected: [0.7310585786300049,0.2689414213699951]

case 2

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

expected: [0.5,0.5]

case 3

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

expected: [0.11920292202211755,0.8807970779778823]

Starter codePython

import math
def trajectory_transformer_attention(query, keys, temperature):
    # Your code here
    pass

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