Apply rotary position embeddings to a block of query vectors and a block of key vectors, where each token sits at its own position in the sequence.
Let D be the head size and half = D / 2. Dimension i of every vector is paired with dimension i + half, for i from 0 to half minus 1. For a token at position p, pair i turns by the angle ai, where
θi=base−2i/D,ai=pθi,and the pair (xi,xi+half) becomes
xi′=xicosai−xi+halfsinai,xi+half′=xi+halfcosai+xisinai.Here q holds query vectors of shape (T, Hq, D) for T tokens and Hq query heads, k holds key vectors of shape (T, Hk, D) for the same tokens, positions holds the position p of each token, and base sets how quickly the turning speeds fall from one pair to the next. Every head of a token turns by the same angles.
Your function is called under jax.jit with q, k and positions traced and base passed as a static Python float.
Return a dict with key q holding the rotated queries and key k holding the rotated keys, both float32 JAX arrays with the same shapes as the inputs.
Input: q = [[[1,0,0,1]],[[1,0,0,1]]], k = [[[0.5,1,-1,2]],[[0.5,1,-1,2]]], positions = [0,1], base = 10000
Output: {"q":[[[1.0,0.0,0.0,1.0]],[[0.540302,-0.01,0.841471,0.99995]]],"k":[[[0.5,1.0,-1.0,2.0]],[[1.111622,0.97995,-0.119567,2.0099]]]}
Explanation: Position 0 has every angle equal to 0, so the first token comes back unchanged. At position 1 with D = 4 and base 10000 the two speeds are 1 and 0.01, so the query pair (q[0], q[2]) = (1, 0) turns by 1 radian into (0.540302, 0.841471), and the pair (q[1], q[3]) = (0, 1) turns by 0.01 radians into (-0.01, 0.99995).
Input: q = [[[1,0],[0,1]],[[2,1],[-1,0.5]]], k = [[[1,1]],[[0,-1]]], positions = [3,7], base = 100
Output: {"q":[[[-0.989992,0.14112],[-0.14112,-0.989992]],[[0.850818,2.067875],[-1.082396,-0.280035]]],"k":[[[-1.131113,-0.848872]],[[0.656987,-0.753902]]]}
Input: q = [[[0.2,-0.4,0.6,1.0,0.0,-0.5]]], k = [[[1,0,0,0,1,0]]], positions = [100], base = 10
Output: {"q":[[[0.678829,0.303865,-0.325058,0.761046,-0.260127,0.710167]]],"k":[[[0.862319,-0.650317,0.0,-0.506366,-0.759663,0.0]]]}
jnp.arange(half, dtype=jnp.float32) gives the pair indices, and base ** array raises base to each one.
angles[:, None, :] broadcasts a (T, half) table across heads.
x[..., :half], x[..., half:] and jnp.concatenate([a, b], axis=-1).
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Apply rotary position embeddings to a block of query vectors and a block of key vectors, where each token sits at its own position in the sequence.
Let D be the head size and half = D / 2. Dimension i of every vector is paired with dimension i + half, for i from 0 to half minus 1. For a token at position p, pair i turns by the angle ai, where
θi=base−2i/D,ai=pθi,and the pair (xi,xi+half) becomes
xi′=xicosai−xi+halfsinai,xi+half′=xi+halfcosai+xisinai.Here q holds query vectors of shape (T, Hq, D) for T tokens and Hq query heads, k holds key vectors of shape (T, Hk, D) for the same tokens, positions holds the position p of each token, and base sets how quickly the turning speeds fall from one pair to the next. Every head of a token turns by the same angles.
Your function is called under jax.jit with q, k and positions traced and base passed as a static Python float.
Return a dict with key q holding the rotated queries and key k holding the rotated keys, both float32 JAX arrays with the same shapes as the inputs.
Input: q = [[[1,0,0,1]],[[1,0,0,1]]], k = [[[0.5,1,-1,2]],[[0.5,1,-1,2]]], positions = [0,1], base = 10000
Output: {"q":[[[1.0,0.0,0.0,1.0]],[[0.540302,-0.01,0.841471,0.99995]]],"k":[[[0.5,1.0,-1.0,2.0]],[[1.111622,0.97995,-0.119567,2.0099]]]}
Explanation: Position 0 has every angle equal to 0, so the first token comes back unchanged. At position 1 with D = 4 and base 10000 the two speeds are 1 and 0.01, so the query pair (q[0], q[2]) = (1, 0) turns by 1 radian into (0.540302, 0.841471), and the pair (q[1], q[3]) = (0, 1) turns by 0.01 radians into (-0.01, 0.99995).
Input: q = [[[1,0],[0,1]],[[2,1],[-1,0.5]]], k = [[[1,1]],[[0,-1]]], positions = [3,7], base = 100
Output: {"q":[[[-0.989992,0.14112],[-0.14112,-0.989992]],[[0.850818,2.067875],[-1.082396,-0.280035]]],"k":[[[-1.131113,-0.848872]],[[0.656987,-0.753902]]]}
Input: q = [[[0.2,-0.4,0.6,1.0,0.0,-0.5]]], k = [[[1,0,0,0,1,0]]], positions = [100], base = 10
Output: {"q":[[[0.678829,0.303865,-0.325058,0.761046,-0.260127,0.710167]]],"k":[[[0.862319,-0.650317,0.0,-0.506366,-0.759663,0.0]]]}
jnp.arange(half, dtype=jnp.float32) gives the pair indices, and base ** array raises base to each one.
angles[:, None, :] broadcasts a (T, half) table across heads.
x[..., :half], x[..., half:] and jnp.concatenate([a, b], axis=-1).
Accepts: array
Accepts: array
Accepts: array
Accepts: number
