class_docstr: ## Relative Multi-Head Attention Module
  
  We override [Multi-Head Attention](mha.html) module so we only need to 
  write the `get_scores` method.
docstr: ### Get relative attention scores
  
  With absolute attention
  
  \begin{align}
  A^{abs}_{j} &= lin_q(X^q_i + P_i)^\top lin_k(X^k_j + P_j) \\
                &= \underset{\textcolor{lightgreen}{A}}{Q_i^\top K_j} +
                   \underset{\textcolor{lightgreen}{B}}{Q_i^\top U^K_j} +
                   \underset{\textcolor{lightgreen}{C}}{{U^Q_i}^\top K_j} +
                   \underset{\textcolor{lightgreen}{D}}{{U^Q_i}^\top U^K_j}
  \end{align}
  
  where $Q_i, K_j$, are linear transformations of
   original embeddings $X^q_i, X^k_j$
   and $U^Q_i, U^K_j$ are linear transformations of
   absolute positional encodings $P_i, P_j$.
  
  They reason out that the attention to a given key should be the same regardless of
  the position of query.
  Hence replace $\underset{\textcolor{lightgreen}{C}}{{U^Q_i}^\top K_j}$
  with a constant $\underset{\textcolor{lightgreen}{C}}{\textcolor{orange}{v^\top} K_j}$.
  
  For the second and third terms relative positional encodings are introduced.
  So $\underset{\textcolor{lightgreen}{B}}{Q_i^\top U^K_j}$ is
  replaced with $\underset{\textcolor{lightgreen}{B}}{Q_i^\top \textcolor{orange}{R_{i - j}}}$
  and $\underset{\textcolor{lightgreen}{D}}{{U^Q_i}^\top U^K_j}$
  with $\underset{\textcolor{lightgreen}{D}}{\textcolor{orange}{S_{i-j}}}$.
  
  \begin{align}
  A^{rel}_{i,j} &= \underset{\mathbf{\textcolor{lightgreen}{A}}}{Q_i^\top K_j} +
                   \underset{\mathbf{\textcolor{lightgreen}{B}}}{Q_i^\top \textcolor{orange}{R_{i - j}}} +
                   \underset{\mathbf{\textcolor{lightgreen}{C}}}{\textcolor{orange}{v^\top} K_j} +
                   \underset{\mathbf{\textcolor{lightgreen}{D}}}{\textcolor{orange}{S_{i-j}}}
  \end{align}
code: def get_scores(self, query: torch.Tensor, key: torch.Tensor):
      r"""
          ### Get relative attention scores
  
          With absolute attention
  
          \begin{align}
          A^{abs}_{j} &= lin_q(X^q_i + P_i)^\top lin_k(X^k_j + P_j) \\
                        &= \underset{\textcolor{lightgreen}{A}}{Q_i^\top K_j} +
                           \underset{\textcolor{lightgreen}{B}}{Q_i^\top U^K_j} +
                           \underset{\textcolor{lightgreen}{C}}{{U^Q_i}^\top K_j} +
                           \underset{\textcolor{lightgreen}{D}}{{U^Q_i}^\top U^K_j}
          \end{align}
  
          where $Q_i, K_j$, are linear transformations of
           original embeddings $X^q_i, X^k_j$
           and $U^Q_i, U^K_j$ are linear transformations of
           absolute positional encodings $P_i, P_j$.
  
          They reason out that the attention to a given key should be the same regardless of
          the position of query.
          Hence replace $\underset{\textcolor{lightgreen}{C}}{{U^Q_i}^\top K_j}$
          with a constant $\underset{\textcolor{lightgreen}{C}}{\textcolor{orange}{v^\top} K_j}$.
  
          For the second and third terms relative positional encodings are introduced.
          So $\underset{\textcolor{lightgreen}{B}}{Q_i^\top U^K_j}$ is
          replaced with $\underset{\textcolor{lightgreen}{B}}{Q_i^\top \textcolor{orange}{R_{i - j}}}$
          and $\underset{\textcolor{lightgreen}{D}}{{U^Q_i}^\top U^K_j}$
          with $\underset{\textcolor{lightgreen}{D}}{\textcolor{orange}{S_{i-j}}}$.
  
          \begin{align}
          A^{rel}_{i,j} &= \underset{\mathbf{\textcolor{lightgreen}{A}}}{Q_i^\top K_j} +
                           \underset{\mathbf{\textcolor{lightgreen}{B}}}{Q_i^\top \textcolor{orange}{R_{i - j}}} +
                           \underset{\mathbf{\textcolor{lightgreen}{C}}}{\textcolor{orange}{v^\top} K_j} +
                           \underset{\mathbf{\textcolor{lightgreen}{D}}}{\textcolor{orange}{S_{i-j}}}
          \end{align}
          """
  
      # $\textcolor{orange}{R_k}$
      key_pos_emb = self.key_pos_embeddings[
          self.P - key.shape[0] : self.P + query.shape[0]
      ]
      # $\textcolor{orange}{S_k}$
      key_pos_bias = self.key_pos_bias[
          self.P - key.shape[0] : self.P + query.shape[0]
      ]
      # $\textcolor{orange}{v^\top}$
      query_pos_bias = self.query_pos_bias[None, None, :, :]
  
      # ${(\textcolor{lightgreen}{\mathbf{A + C}})}_{i,j} =
      # Q_i^\top K_j +
      # \textcolor{orange}{v^\top} K_jZ$
      ac = torch.einsum("ibhd,jbhd->ijbh", query + query_pos_bias, key)
      # $\textcolor{lightgreen}{\mathbf{B'}_{i,k}} = Q_i^\top \textcolor{orange}{R_k}$
      b = torch.einsum("ibhd,jhd->ijbh", query, key_pos_emb)
      # $\textcolor{lightgreen}{\mathbf{D'}_{i,k}} = \textcolor{orange}{S_k}$
      d = key_pos_bias[None, :, None, :]
      # Shift the rows of $\textcolor{lightgreen}{\mathbf{(B' + D')}_{i,k}}$
      # to get $$\textcolor{lightgreen}{\mathbf{(B + D)}_{i,j} = \mathbf{(B' + D')}_{i,i - j}}$$
      bd = shift_right(b + d)
      # Remove extra positions
      bd = bd[:, -key.shape[0] :]
  
      # Return the sum $$
      # \underset{\mathbf{\textcolor{lightgreen}{A}}}{Q_i^\top K_j} +
      # \underset{\mathbf{\textcolor{lightgreen}{B}}}{Q_i^\top \textcolor{orange}{R_{i - j}}} +
      # \underset{\mathbf{\textcolor{lightgreen}{C}}}{\textcolor{orange}{v^\top} K_j} +
      # \underset{\mathbf{\textcolor{lightgreen}{D}}}{\textcolor{orange}{S_{i-j}}}
      # $$
      return ac + bd
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/xl/relative_mha.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 RelativeMultiHeadAttention obj's underlying member method code with example usage: obj.get_scores(query, key) and returns: unknown