class_docstr: ## [RoPE embeddings](../rope/index.html)
  
  *We use rotary position embeddings in self-attention layers.
  We assume the positional information gets embedded in embeddings
  and therefore not use them in causal attention.
  [Non-causal self-attention needs explicit positional information
   because it cannot infer it](https://arxiv.org/abs/3999902edc8511eba3db37f65e372566).*
class_docstr_tok: ["#","#","[","RoPE","embeddings","]","(","..","\/rope\/index.html",")","*","We","use","rotary","position","embeddings","in","self-attention","layers",".","We","assume","the","positional","information","gets","embedded","in","embeddings","and","therefore","not","use","them","in","causal","attention",".","[","Non-causal","self-attention","needs","explicit","positional","information","because","it","can","not","infer","it","]","(","https",":","\/\/arxiv.org\/abs\/3999902edc8511eba3db37f65e372566",")",".","*"]
docstr: * `x` is the Tensor at the head of a key or a query with shape `[ batch_size, seq_len, n_heads, d]`
docstr_tok: ["*","`","x","`","is","the","Tensor","at","the","head","of","a","key","or","a","query","with","shape","`","[","batch_size",",","seq_len",",","n_heads",",","d","]","`"]
code: def forward(self, x: torch.Tensor):
      """
      * `x` is the Tensor at the head of a key or a query with shape `[ batch_size, seq_len, n_heads, d]`
      """
      # Extract the shape
      batch_size, seq_len, n_heads, d = x.shape
  
      # $\frac{d}{2}$
      d_2 = d // 2
  
      # Create position indexes `[0, 1, ..., seq_len - 1]`
      seq_idx = torch.arange(seq_len, device=x.device).type_as(
          self.theta
      )
  
      # Calculate the product of position index and $\theta_i$
      idx_theta = torch.einsum("n,d->nd", 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)
  
      # Calculate
      # $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., -x^{(\frac{d}{2})}]$
      neg_half_x = torch.cat(
          [-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1
      )
  
      # Calculate
      #
      # \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}}$
      rx = (x * idx_theta2.cos()[None, :, None, :]) + (
          neg_half_x * idx_theta2.sin()[None, :, None, :]
      )
  
      #
      return rx
  
imports: import math
  from typing import Set
  import torch
  from torch import nn
  from labml.logger import inspect
label_desc: Use image node_id 2 for calling the RotaryPositionalEmbeddings obj's underlying member method code with example usage: obj.forward(x) and returns: rx