Position Interpolation › A Complete Implementation

class RoPEWithPositionInterpolation:
    """RoPE implementation with Position Interpolation support."""

    def __init__(
        self, d_model, base=10000, train_length=2048, target_length=None
    ):
        """Initialize RoPE with optional Position Interpolation.

        Args:
            d_model: Embedding dimension (must be even)
            base: Base for frequency computation
            train_length: Maximum position seen during training
            target_length: Extended context length (None = no interpolation)
        """
        if d_model % 2 != 0:
            raise ValueError("d_model must be even for RoPE")

        self.d_model = d_model
        self.base = base
        self.train_length = train_length

        # Compute scale factor for Position Interpolation
        if target_length is not None and target_length > train_length:
            self.scale = train_length / target_length
        else:
            self.scale = 1.0

        # Precompute frequencies
        dim_pairs = d_model // 2
        i = np.arange(dim_pairs)
        self.theta = 1.0 / (base ** (2 * i / d_model))

    def get_rotation_matrix(self, position):
        """Get the block-diagonal rotation matrix for a position.

        Args:
            position: Integer position index

        Returns:
            R: Rotation matrix of shape (d_model, d_model)
        """
        # Apply position scaling
        scaled_pos = position * self.scale

        # Compute angles for each dimension pair
        angles = scaled_pos * self.theta

        # Build block-diagonal rotation matrix
        R = np.zeros((self.d_model, self.d_model))
        for i, angle in enumerate(angles):
            cos_a, sin_a = np.cos(angle), np.sin(angle)
            idx = 2 * i
            R[idx, idx] = cos_a
            R[idx, idx + 1] = -sin_a
            R[idx + 1, idx] = sin_a
            R[idx + 1, idx + 1] = cos_a

        return R

    def apply(self, x, positions):
        """Apply RoPE to input vectors.

        Args:
            x: Input array of shape (seq_len, d_model)
            positions: Position indices for each element

        Returns:
            rotated: Rotated vectors of shape (seq_len, d_model)
        """
        seq_len = x.shape[0]
        rotated = np.zeros_like(x)

        for i, pos in enumerate(positions):
            R = self.get_rotation_matrix(pos)
            rotated[i] = R @ x[i]

        return rotated

# Create instances with and without interpolation
rope_standard = RoPEWithPositionInterpolation(d_model=64, train_length=2048)

rope_interpolated = RoPEWithPositionInterpolation(
    d_model=64, train_length=2048, target_length=8192
)

# Test on a sample vector at an extended position. The hidden setup cell seeds
# the chapter once so all render variants stay deterministic.
test_vector = np.random.randn(64)
test_position = 6000  # Beyond training range

# Get rotation matrices
R_standard = rope_standard.get_rotation_matrix(test_position)
R_interpolated = rope_interpolated.get_rotation_matrix(test_position)
R_training_equiv = rope_standard.get_rotation_matrix(test_position * 0.25)