class_docstr: ## Feedback Attention
  
  
  This module computes recurrent attention similar to attention from original transformers
  paper.
  
  $$\mathop{Attention}(Q, K, V) = \underset{seq}{\mathop{softmax}}\Bigg(\frac{Q^\top K}{\sqrt{d_k}}\Bigg)V$$
class_docstr_tok: ["#","#","Feedback","Attention","This","module","computes","recurrent","attention","similar","to","attention","from","original","transformers","paper",".","$","$","\\mathop","{","Attention","}","(","Q",",","K",",","V",")","=","\\underset","{","seq","}","{","\\mathop","{","softmax","}","}","\\Bigg","(","\\frac","{","Q^\\top","K","}","{","\\sqrt","{","d_k","}","}","\\Bigg",")","V","$","$"]
docstr: ### Get attention scores
  
  We use relative positional encodings for attention, similar
  to [relative multi-head attention form Transformer-XL paper](../relative_mha.html).
  
  Attention from current step's query to key in step $j$ (relative to current step) is,
  
  \begin{align}
  A_{j} &= Q^\top K_j \\
      &= lin_q(X^q + P_q)^\top lin_k(X^k_j + P_j) \\
      &= (Q + U^Q)^\top(K_j + U^K_j) \\
      &= \underset{\textcolor{lightgreen}{A}}{Q^\top K_j} +
         \underset{\textcolor{lightgreen}{B}}{Q^\top U^K_j} +
         \underset{\textcolor{lightgreen}{C}}{{U^Q}^\top K_j} +
         \underset{\textcolor{lightgreen}{D}}{{U^Q}^\top U^K_j}
  \end{align}
  
  where $Q, K_j$, are linear transformations of
   original embeddings $X^q, X^k_j$
   and $U^Q, U^K_j$ are linear transformations of
   positional encodings $P_q, P_j$.
  
  We replace term $\textcolor{lightgreen}{D}$ with $S_j$.
docstr_tok: ["#","#","#","Get","attention","scores","We","use","relative","positional","encodings","for","attention",",","similar","to","[","relative","multi-head","attention","form","Transformer-XL","paper","]","(","..","\/relative_mha.html",")",".","Attention","from","current","step","'s","query","to","key","in","step","$","j","$","(","relative","to","current","step",")","is",",","\\begin","{","align","}","A_","{","j","}","&","=","Q^\\top","K_j","\\\\","&","=","lin_q","(","X^q","+","P_q",")","^\\top","lin_k","(","X^k_j","+","P_j",")","\\\\","&","=","(","Q","+","U^Q",")","^\\top","(","K_j","+","U^K_j",")","\\\\","&","=","\\underset","{","\\textcolor","{","lightgreen","}","{","A","}","}","{","Q^\\top","K_j","}","+","\\underset","{","\\textcolor","{","lightgreen","}","{","B","}","}","{","Q^\\top","U^K_j","}","+","\\underset","{","\\textcolor","{","lightgreen","}","{","C","}","}","{","{","U^Q","}","^\\top","K_j","}","+","\\underset","{","\\textcolor","{","lightgreen","}","{","D","}","}","{","{","U^Q","}","^\\top","U^K_j","}","\\end","{","align","}","where","$","Q",",","K_j","$",",","are","linear","transformations","of","original","embeddings","$","X^q",",","X^k_j","$","and","$","U^Q",",","U^K_j","$","are","linear","transformations","of","positional","encodings","$","P_q",",","P_j","$",".","We","replace","term","$","\\textcolor","{","lightgreen","}","{","D","}","$","with","$","S_j","$","."]
code: def get_scores(self, query: torch.Tensor, key: torch.Tensor):
      r"""
          ### Get attention scores
  
          We use relative positional encodings for attention, similar
          to [relative multi-head attention form Transformer-XL paper](../relative_mha.html).
  
          Attention from current step's query to key in step $j$ (relative to current step) is,
  
          \begin{align}
          A_{j} &= Q^\top K_j \\
              &= lin_q(X^q + P_q)^\top lin_k(X^k_j + P_j) \\
              &= (Q + U^Q)^\top(K_j + U^K_j) \\
              &= \underset{\textcolor{lightgreen}{A}}{Q^\top K_j} +
                 \underset{\textcolor{lightgreen}{B}}{Q^\top U^K_j} +
                 \underset{\textcolor{lightgreen}{C}}{{U^Q}^\top K_j} +
                 \underset{\textcolor{lightgreen}{D}}{{U^Q}^\top U^K_j}
          \end{align}
  
          where $Q, K_j$, are linear transformations of
           original embeddings $X^q, X^k_j$
           and $U^Q, U^K_j$ are linear transformations of
           positional encodings $P_q, P_j$.
  
          We replace term $\textcolor{lightgreen}{D}$ with $S_j$.
          """
  
      # $U^K_j$
      key_pos_emb = self.key_pos_embeddings[-key.shape[0] :]
      # $U^Q$
      query_pos_bias = self.query_pos_bias[None, :, :]
      # $S_j$
      key_pos_bias = self.key_pos_bias[-key.shape[0] :]
  
      # $\underset{\textcolor{lightgreen}{A}}{Q^\top K_j} + \underset{\textcolor{lightgreen}{C}}{{U^Q}^\top K_j}$
      ac = torch.einsum("bhd,jbhd->jbh", query + query_pos_bias, key)
      # $\underset{\textcolor{lightgreen}{B}}{Q^\top U^K_j} + \underset{\textcolor{lightgreen}{D}}{S_j}$
      bd = (
          torch.einsum("bhd,jhd->jbh", query, key_pos_emb)
          + key_pos_bias[:, None, :]
      )
  
      # $A_j$
      return ac + bd
  
imports: import math
  from typing import Optional
  import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.transformers.feed_forward import FeedForward
  from labml_nn.transformers.mha import PrepareForMultiHeadAttention
  from labml_nn.utils import clone_module_list
label_desc: Use image node_id 2 for calling the FeedbackAttention obj's underlying member method code with example usage: obj.get_scores(query, key) and returns: unknown