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: * 'heads' is the number of attention heads
  * `d_model` is the number of features in the transformer
  * `dropout_prob` is the attention dropout probability
  * `is_kv_precomputed` is whether key, value tensors are already calculated
docstr_tok: ["*","'heads","'","is","the","number","of","attention","heads","*","`","d_model","`","is","the","number","of","features","in","the","transformer","*","`","dropout_prob","`","is","the","attention","dropout","probability","*","`","is_kv_precomputed","`","is","whether","key",",","value","tensors","are","already","calculated"]
code: def __init__(
      self,
      heads: int,
      d_model: int,
      dropout_prob: float = 0.1,
      *,
      is_kv_precomputed: bool = False,
  ):
      """
      * 'heads' is the number of attention heads
      * `d_model` is the number of features in the transformer
      * `dropout_prob` is the attention dropout probability
      * `is_kv_precomputed` is whether key, value tensors are already calculated
      """
  
      super().__init__()
  
      # Number of features per head
      self.d_k = d_model // heads
      #
      self.heads = heads
  
      # These transform the `query` multi-headed attention.
      self.query = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=False
      )
      # These transform the `key` and `value` for multi-headed attention.
      if not is_kv_precomputed:
          self.key = PrepareForMultiHeadAttention(
              d_model, heads, self.d_k, bias=False
          )
          self.value = PrepareForMultiHeadAttention(
              d_model, heads, self.d_k, bias=True
          )
      # Keys and values are already calculated
      else:
          self.key = None
          self.value = None
  
      # Output layer
      self.output = nn.Linear(d_model, d_model)
      # Dropout
      self.dropout = nn.Dropout(dropout_prob)
      # Scaling factor before the softmax
      self.scale = 1 / math.sqrt(self.d_k)
  
      # Softmax for attention along the time dimension of `key`
      self.softmax = nn.Softmax(dim=0)
  
      # Number of relative positions
      self.P = 2**12
  
      # Relative positional embeddings for key relative to the query.
      self.key_pos_embeddings = nn.Parameter(
          torch.zeros((self.P, heads, self.d_k)), requires_grad=True
      )
      # Relative positional embedding bias for key relative to the query.
      self.key_pos_bias = nn.Parameter(
          torch.zeros((self.P, heads)), requires_grad=True
      )
      # Positional embeddings for the query is independent of the position of the query
      self.query_pos_bias = nn.Parameter(
          torch.zeros((heads, self.d_k)), requires_grad=True
      )
  
      # We store attentions so that it can be used for logging, or other computations if needed
      self.attn = None
  
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 1 to create a new FeedbackAttention object from inherited base classes: Module with example: obj = FeedbackAttention(heads, d_model, dropout_prob)