RoPE Relative Position Dot

Problem

Apply rotary embeddings to q at position m and k at position n (same even dimension, angle position * theta_base^(-2p/d) for pair p), then return their dot product. The result depends only on m - n.

The signature is rope_relative_dot(q, k, m, n, theta_base=10000.0).

Hint

Rotating both vectors by their positions makes the attention score a function of the offset m - n.

Test cases

case 1

in: [[1,0],[1,0],5,5,10000]

expected: 1

case 2

in: [[1,0],[1,0],3,1,10000]

expected: -0.4161468365

case 3

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

expected: 0

case 4

in: [[1,0,1,0],[1,0,1,0],7,4,10000]

expected: 0.0095575371

case 5

in: [[1,2,3,4],[4,3,2,1],2,2,10000]

expected: 20

Starter codePython

import math
def rope_relative_dot(q, k, m, n, theta_base=10000.0):
    # Your code here
    pass

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