Project: s41467-025-60872-5 | Function: apply_rotary_pos_emb (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements Rotary Position Embedding (RoPE) algorithm for encoding positional information in transformer attention. Applies rotation transformations to query and key vectors using precomputed cosine and sine values, enabling better handling of relative positions in RNA sequences.

**KEY TERMS:** {
  "q": "query vectors",
  "k": "key vectors",
  "cos": "cosine positional encoding",
  "sin": "sine positional encoding",
  "rotate_half": "rotation operation for embedding"
}

**CONTEXT CODE:** import torch
import torch.nn as nn

# Code heavily inspired by https://blog.eleuther.ai/rotary-embeddings/, GPT-NeoX (Pytorch) implementation
# and ESM2 implementation https://github.com/facebookresearch/esm/blob/main/esm/rotary_embedding.py

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    # MASKED: Implements Rotary Position Embedding (RoPE) algorithm for encoding positional information using rotation transformations on query and key vectors with precomputed cosine and sine values implementation
    raise NotImplementedError("This function is masked for testing purposes.")

class RotaryPositionEmbedding(nn.Module):
    def __init__(
        self,
        dim: int,
        base: int = 10000
    ):
        super().__init__()

        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

        self.seq_len_cached = None
        self.cos_cached = None
        self.sin_cached = None

    def _update_cached(self, x, seq_dim):
        seq_len = x.shape[seq_dim]

        if seq_len != self.seq_len_cached:
            self.seq_len_cached = seq_len
            t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)

            self.cos_cached = emb.cos()[None, None, :, :]
            self.sin_cached = emb.sin()[None, None, :, :]

    def forward(self, q, k):
        self._update_cached(k, seq_dim=-2)
        return apply_rotary_pos_emb(q, k, self.cos_cached, self.sin_cached)


Complete the function implementation replacing "raise NotImplementedError".