Skip to content

← Position Encodings step 3 of 5

Hard Research

Rotary Position Embeddings

Rotary position embeddings (RoPE) encode position by rotating feature pairs. Think of two adjacent features as the coordinates of a point in a plane. The position chooses how far to rotate that point; the pair's frequency controls how quickly the angle grows.

Your task

Implement apply_rope(x, positions).

  • x is a floating-point array of shape (seq_len, d), with positive even d.
  • positions has shape (seq_len,) and contains an integer position for each row.
  • seq_len is positive. Use the supplied positions, not the row indices.

Return shape (seq_len, d) without changing either input. Preserve x's floating-point dtype and, in PyTorch, its device. Positions need not start at zero or be consecutive; repeated positions use the same rotation angles.

This exercise pairs adjacent columns (0, 1), (2, 3), and so on. For pair index i and row s, compute:

frequency = 10000**(-2 * i / d)
angle = positions[s] * frequency
a = x[s, 2 * i]
b = x[s, 2 * i + 1]
output[s, 2 * i]     = a * cos(angle) - b * sin(angle)
output[s, 2 * i + 1] = a * sin(angle) + b * cos(angle)

Use radians and ordinary division in the exponent. Both output coordinates must use the original a and b. Do not use an already-updated coordinate to compute the second one. The pair index i ranges from zero to d // 2 - 1.

Example

For x = [[1, 0, 0, 1]] at position 1, the two frequencies are 1 and 0.01. The first pair rotates [1, 0] by one radian; the second rotates [0, 1] by 0.01 radians. The result is approximately:

[[0.540302, 0.841471, -0.010000, 0.999950]]

At position zero every angle is zero, so the row stays unchanged. A rotation also preserves each pair's squared length, up to floating-point rounding: a*a + b*b equals the sum of squares of the rotated coordinates.

This function only rotates the supplied vectors. It does not add an encoding table, project queries or keys, compute attention, or rotate pairs made from the first and second halves of the feature axis.

Keep angles as array values rather than converting positions to Python floats. In PyTorch, compute trigonometric values on x's device with its dtype.

The visualization uses four pairs (d = 8), each starting at [1, 0]. Move the position slider to compare their rotation rates. Your inputs may start with different lengths and directions, as the example above does.

Loading visualization…