Rotary Embedding Apply 2D

Problem

Apply rotary position embeddings (RoPE) to a vector of even dimension. For each adjacent pair i, i+1 with pair index p = i // 2, the rotation angle is position * theta_base^(-2p/d), and the pair becomes [x0*cos - x1*sin, x0*sin + x1*cos].

The signature is rotary_embedding_apply(x, position, theta_base=10000.0).

Hint

Higher pair indices rotate at lower frequencies, encoding different position scales.

Test cases

case 1

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

expected: [1,0]

case 2

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

expected: [0.5403023059,0.8414709848]

case 3

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

expected: [0.5403023059,0.8414709848,0.9999500004,0.0099998333]

case 4

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

expected: [-0.4161468365,0.9092974268,-0.4161468365,0.9092974268]

Starter codePython

import math
def rotary_embedding_apply(x, position, theta_base=10000.0):
    # Your code here
    pass

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