class_docstr: ## Rotary Positional Embeddings
  
  GPT-NeoX uses [rotary positional embeddings (RoPE)](https://arxiv.org/abs/2104.09864).
  
  WE have annotated implementation of RoPE [here](https://nn.labml.ai/transformers/rope/index.html)
  with more notes the theory.
class_docstr_tok: ["#","#","Rotary","Positional","Embeddings","GPT-NeoX","uses","[","rotary","positional","embeddings","(","RoPE",")","]","(","https",":","\/\/arxiv.org\/abs\/2104.09864",")",".","WE","have","annotated","implementation","of","RoPE","[","here","]","(","https",":","\/\/nn.labml.ai\/transformers\/rope\/index.html",")","with","more","notes","the","theory","."]
docstr: :param x: has shape `[..., seq, n_heads, d_k]`
  :param offset: is the starting position of `x`. This is $\gt 0$ when we have
  cached the keys and queries of previous positions
code: def forward(self, x: torch.Tensor, offset: int = 0):
      """
      :param x: has shape `[..., seq, n_heads, d_k]`
      :param offset: is the starting position of `x`. This is $\gt 0$ when we have
      cached the keys and queries of previous positions
      """
  
      # Get the actual sequence length
      seq_len = x.shape[-3] + offset
  
      # Initialize $\theta$
      if self.theta is None:
          #  $\theta_i = 10000^{\frac{2(i-1)}{d}}$
          theta = 1.0 / (
              self.base
              ** (torch.arange(0, self.d_rope, 2).float() / self.d_rope)
          )
          self.theta = theta.to(x.device).to(x.dtype)
  
      # Initialize $\cos m\theta_i$ and $\sin m\theta_i$ cache
      if (
          self.cos_cached is None
          or seq_len > self.cos_cached.shape[1]
          or self.cos_cached.device != x.device
          or self.cos_cached.dtype != x.dtype
      ):
          # Get position indexes $m$
          seq_idx = torch.arange(seq_len, device=x.device).type_as(
              self.theta
          )
          # $m \theta_i$
          idx_theta = torch.einsum("s,d->sd", seq_idx, self.theta)
          # Concatenate so that for row $m$ we have
          #
          # $$[m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}, m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}]$$
          idx_theta2 = torch.cat((idx_theta, idx_theta), dim=-1).to(
              x.device
          )
  
          # Calculate $\cos m\theta_i$ and $\sin m\theta_i$ in fp32
          with autocast(enabled=False):
              idx_theta2 = idx_theta2.float()
              # Add head dimension
              self.cos_cached = idx_theta2.cos()[:, None, :]
              self.sin_cached = idx_theta2.sin()[:, None, :]
  
          # Cache them
          self.cos_cached = self.cos_cached.to(x.dtype)
          self.sin_cached = self.sin_cached.to(x.dtype)
  
      # Split the features. We apply RoPE to only `d_rope` features
      x_rope, x_pass = x[..., : self.d_rope], x[..., self.d_rope :]
  
      # Get the sin and cos values from the cache
      cos, sin = (
          self.cos_cached[offset:seq_len],
          self.sin_cached[offset:seq_len],
      )
  
      # RoPE embeddings
      #
      # \begin{align}
      # \begin{pmatrix}
      # x^{(i)}_m \cos m \theta_i - x^{(i + \frac{d}{2})}_m \sin m \theta_i \\
      # x^{(i + \frac{d}{2})}_m \cos m\theta_i + x^{(i)}_m \sin m \theta_i \\
      # \end{pmatrix} \\
      # \end{align}
      #
      # for $i \in {1, 2, ..., \frac{d}{2}}$
      x_rope = (x_rope * cos) + (self.rotate_half(x_rope) * sin)
  
      # Concatenate with features that didn't get RoPE embeddings
      return torch.cat((x_rope, x_pass), dim=-1)
  
imports: import copy
  import math
  from typing import Dict, Optional, Set, Callable, Any, Generator, Tuple
  import torch
  from torch import nn
  from torch.cuda.amp import autocast
  from labml import monit, logger
  from labml.logger import Text
  from labml_nn.neox import checkpoint
  from labml_nn.neox.utils.cache import get_cache
label_desc: Use image node_id 3 for calling the RoPE obj's underlying member method code with example usage: obj.forward(x, offset) and returns: torch