class_docstr: ## RoPE module
  
  Rotary encoding transforms pairs of features by rotating in the 2D plane.
  That is, it organizes the $d$ features as $rac{d}{2}$ pairs.
  Each pair can be considered a coordinate in a 2D plane, and the encoding will rotate it
  by an angle depending on the position of the token.
  
  ### For a pair of features
  
  Let $x^{(1)}_m$ and $x^{(2)}_m$ be two features of the
  key or query of any head at position $m$.
  Or for simplicity assume $x$ has only two features.
  Then the transformation is,
  
  egin{align}
  RoPEig(x^{(1)}_m, x^{(2)}_m, mig) &=
  egin{pmatrix}
  \cos m      heta & - \sin m         heta \
  \sin m      heta & \cos m   heta
  \end{pmatrix}
  egin{pmatrix}
  x^{(1)}_m \
  x^{(2)}_m \
  \end{pmatrix} \
  &=
  egin{pmatrix}
  x^{(1)}_m \cos m    heta - x^{(2)}_m \sin m         heta \
  x^{(2)}_m \cos m    heta + x^{(1)}_m \sin m         heta \
  \end{pmatrix} \
  \end{align}
  
  where $     heta$ is a constant angle. The other pairs of features are transformed similarly.
  
  ### Attention is relative
  
  For a pair of features, dot-product attention score between two positions $m$ and $n$ would be
  
  egin{align}
  \Big \langle RoPEig(x^{(1)}_m, x^{(2)}_m, mig),  RoPEig(x^{(1)}_n, x^{(2)}_n, nig) \Big 
  angle &= \
  (x^{(1)}_m \cos m   heta - x^{(2)}_m \sin m         heta)(x^{(1)}_n \cos n  heta - x^{(2)}_n \sin n         heta) &+ \
  (x^{(2)}_m \cos m   heta + x^{(1)}_m \sin m         heta)(x^{(2)}_n \cos n  heta + x^{(1)}_n \sin n         heta) &= \
  x^{(1)}_m x^{(1)}_n (\cos m heta \cos n     heta + \sin m   heta \sin n     heta) &+ \
  x^{(1)}_m x^{(2)}_n (-\cos m        heta \sin n     heta + \sin m   heta \cos n     heta) &+ \
  x^{(2)}_m x^{(1)}_n (-\sin m        heta \cos n     heta + \cos m   heta \sin n     heta) &+ \
  x^{(2)}_m x^{(2)}_n (\sin m heta \sin n     heta + \cos m   heta \cos n     heta) &= \
  
  x^{(1)}_m x^{(1)}_n \cos (m - n)    heta +
  x^{(1)}_m x^{(2)}_n \sin(m - n)     heta &+ \
  - x^{(2)}_m x^{(1)}_n \sin (m - n)  heta +
  x^{(2)}_m x^{(2)}_n \cos (m - n)    heta &= \
  
  ig(x^{(1)}_m \cos (m - n)  heta - x^{(2)}_m \sin (m - n)   hetaig) x^{(1)}_n &+ \
  ig(x^{(2)}_m \cos (m - n)m heta + x^{(1)}_m \sin (m - n)   hetaig) x^{(2)}_n  &= \
  
  \Big \langle RoPEig(x^{(1)}_m, x^{(2)}_m, m - nig),  RoPEig(x^{(1)}_n, x^{(2)}_n, 0ig) \Big 
  angle
  \end{align}
  
  This shows that for dot-production attention the rotary encodings gives relative attention.
  
  ### For all features
  
  The features are grouped into pairs and handled as above. They use a different $    heta$ for each pair.
  
  The paper suggests using $\Theta = {        heta_i = 10000^{rac{2(i-1)}{d}}, i \in [1, 2, ..., rac{d}{2}]}$
  for the $rac{d}{2}$ pairs of features.
  
  We pair feature $i$ with feature $i + rac{d}{2}$. So for position $m$ we transform
  
  egin{align}
  egin{pmatrix}
  x^{(i)}_m \
  x^{(i + rac{d}{2})}_m
  \end{pmatrix}
  \end{align}
  
  to
  
  egin{align}
  egin{pmatrix}
  x^{(i)}_m \cos m    heta_i - x^{(i + rac{d}{2})}_m \sin m  heta_i \
  x^{(i + rac{d}{2})}_m \cos m       heta_i + x^{(i)}_m \sin m       heta_i \
  \end{pmatrix} \
  \end{align}
class_docstr_tok: ["#","#","RoPE","module","Rotary","encoding","transforms","pairs","of","features","by","rotating","in","the","2D","plane",".","That","is",",","it","organizes","the","$","d","$","features","as","$","rac","{","d","}","{","2","}","$","pairs",".","Each","pair","can","be","considered","a","coordinate","in","a","2D","plane",",","and","the","encoding","will","rotate","it","by","an","angle","depending","on","the","position","of","the","token",".","#","#","#","For","a","pair","of","features","Let","$","x^","{","(","1",")","}","_m","$","and","$","x^","{","(","2",")","}","_m","$","be","two","features","of","the","key","or","query","of","any","head","at","position","$","m","$",".","Or","for","simplicity","assume","$","x","$","has","only","two","features",".","Then","the","transformation","is",",","\begin","{","align","}","RoPE\big","(","x^","{","(","1",")","}","_m",",","x^","{","(","2",")","}","_m",",","m\big",")","&","=","\begin","{","pmatrix","}","\\cos","m","heta","&","-","\\sin","m","heta","\\","\\sin","m","heta","&","\\cos","m","heta","\\end","{","pmatrix","}","\begin","{","pmatrix","}","x^","{","(","1",")","}","_m","\\","x^","{","(","2",")","}","_m","\\","\\end","{","pmatrix","}","\\","&","=","\begin","{","pmatrix","}","x^","{","(","1",")","}","_m","\\cos","m","heta","-","x^","{","(","2",")","}","_m","\\sin","m","heta","\\","x^","{","(","2",")","}","_m","\\cos","m","heta","+","x^","{","(","1",")","}","_m","\\sin","m","heta","\\","\\end","{","pmatrix","}","\\","\\end","{","align","}","where","$","heta","$","is","a","constant","angle",".","The","other","pairs","of","features","are","transformed","similarly",".","#","#","#","Attention","is","relative","For","a","pair","of","features",",","dot-product","attention","score","between","two","positions","$","m","$","and","$","n","$","would","be","\begin","{","align","}","\\Big","\\langle","RoPE\big","(","x^","{","(","1",")","}","_m",",","x^","{","(","2",")","}","_m",",","m\big",")",",","RoPE\big","(","x^","{","(","1",")","}","_n",",","x^","{","(","2",")","}","_n",",","n\big",")","\\Big","angle","&","=","\\","(","x^","{","(","1",")","}","_m","\\cos","m","heta","-","x^","{","(","2",")","}","_m","\\sin","m","heta",")","(","x^","{","(","1",")","}","_n","\\cos","n","heta","-","x^","{","(","2",")","}","_n","\\sin","n","heta",")","&","+","\\","(","x^","{","(","2",")","}","_m","\\cos","m","heta","+","x^","{","(","1",")","}","_m","\\sin","m","heta",")","(","x^","{","(","2",")","}","_n","\\cos","n","heta","+","x^","{","(","1",")","}","_n","\\sin","n","heta",")","&","=","\\","x^","{","(","1",")","}","_m","x^","{","(","1",")","}","_n","(","\\cos","m","heta","\\cos","n","heta","+","\\sin","m","heta","\\sin","n","heta",")","&","+","\\","x^","{","(","1",")","}","_m","x^","{","(","2",")","}","_n","(","-\\cos","m","heta","\\sin","n","heta","+","\\sin","m","heta","\\cos","n","heta",")","&","+","\\","x^","{","(","2",")","}","_m","x^","{","(","1",")","}","_n","(","-\\sin","m","heta","\\cos","n","heta","+","\\cos","m","heta","\\sin","n","heta",")","&","+","\\","x^","{","(","2",")","}","_m","x^","{","(","2",")","}","_n","(","\\sin","m","heta","\\sin","n","heta","+","\\cos","m","heta","\\cos","n","heta",")","&","=","\\","x^","{","(","1",")","}","_m","x^","{","(","1",")","}","_n","\\cos","(","m","-","n",")","heta","+","x^","{","(","1",")","}","_m","x^","{","(","2",")","}","_n","\\sin","(","m","-","n",")","heta","&","+","\\","-","x^","{","(","2",")","}","_m","x^","{","(","1",")","}","_n","\\sin","(","m","-","n",")","heta","+","x^","{","(","2",")","}","_m","x^","{","(","2",")","}","_n","\\cos","(","m","-","n",")","heta","&","=","\\","\big","(","x^","{","(","1",")","}","_m","\\cos","(","m","-","n",")","heta","-","x^","{","(","2",")","}","_m","\\sin","(","m","-","n",")","heta\big",")","x^","{","(","1",")","}","_n","&","+","\\","\big","(","x^","{","(","2",")","}","_m","\\cos","(","m","-","n",")","m","heta","+","x^","{","(","1",")","}","_m","\\sin","(","m","-","n",")","heta\big",")","x^","{","(","2",")","}","_n","&","=","\\","\\Big","\\langle","RoPE\big","(","x^","{","(","1",")","}","_m",",","x^","{","(","2",")","}","_m",",","m","-","n\big",")",",","RoPE\big","(","x^","{","(","1",")","}","_n",",","x^","{","(","2",")","}","_n",",","0\big",")","\\Big","angle","\\end","{","align","}","This","shows","that","for","dot-production","attention","the","rotary","encodings","gives","relative","attention",".","#","#","#","For","all","features","The","features","are","grouped","into","pairs","and","handled","as","above",".","They","use","a","different","$","heta","$","for","each","pair",".","The","paper","suggests","using","$","\\Theta","=","{","heta_i","=","10000^","{","rac","{","2","(","i-1",")","}","{","d","}","}",",","i","\\in","[","1",",","2",",","...",",","rac","{","d","}","{","2","}","]","}","$","for","the","$","rac","{","d","}","{","2","}","$","pairs","of","features",".","We","pair","feature","$","i","$","with","feature","$","i","+","rac","{","d","}","{","2","}","$",".","So","for","position","$","m","$","we","transform","\begin","{","align","}","\begin","{","pmatrix","}","x^","{","(","i",")","}","_m","\\","x^","{","(","i","+","rac","{","d","}","{","2","}",")","}","_m","\\end","{","pmatrix","}","\\end","{","align","}","to","\begin","{","align","}","\begin","{","pmatrix","}","x^","{","(","i",")","}","_m","\\cos","m","heta_i","-","x^","{","(","i","+","rac","{","d","}","{","2","}",")","}","_m","\\sin","m","heta_i","\\","x^","{","(","i","+","rac","{","d","}","{","2","}",")","}","_m","\\cos","m","heta_i","+","x^","{","(","i",")","}","_m","\\sin","m","heta_i","\\","\\end","{","pmatrix","}","\\","\\end","{","align","}"]
docstr: Cache $\cos$ and $\sin$ values
docstr_tok: ["Cache","$","\\cos","$","and","$","\\sin","$","values"]
code: def _build_cache(self, x: torch.Tensor):
      """
      Cache $\cos$ and $\sin$ values
      """
      # Return if cache is already built
      if (
          self.cos_cached is not None
          and x.shape[0] <= self.cos_cached.shape[0]
      ):
          return
  
      # Get sequence length
      seq_len = x.shape[0]
  
      # $\Theta = {\theta_i = 10000^{-\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
      theta = 1.0 / (
          self.base ** (torch.arange(0, self.d, 2).float() / self.d)
      ).to(x.device)
  
      # Create position indexes `[0, 1, ..., seq_len - 1]`
      seq_idx = (
          torch.arange(seq_len, device=x.device).float().to(x.device)
      )
  
      # Calculate the product of position index and $\theta_i$
      idx_theta = torch.einsum("n,d->nd", seq_idx, 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)
  
      # Cache them
      self.cos_cached = idx_theta2.cos()[:, None, None, :]
      self.sin_cached = idx_theta2.sin()[:, None, None, :]
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/rope/__init__.py
imports: import torch
  from torch import nn
  from labml.logger import inspect
  from labml_nn.transformers.mha import MultiHeadAttention
label_desc: Use image node_id 2 for calling the RotaryPositionalEmbeddings obj's underlying member method code with example usage: obj._build_cache(x) without return types