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: * `query` has shape `[batch_size, d_model]`
  * `key` and `value` has shape `[seq_len, batch_size, d_model]`
docstr_tok: ["*","`","query","`","has","shape","`","[","batch_size",",","d_model","]","`","*","`","key","`","and","`","value","`","has","shape","`","[","seq_len",",","batch_size",",","d_model","]","`"]
code: def forward(
      self,
      *,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
  ):
      """
      * `query` has shape `[batch_size, d_model]`
      * `key` and `value` has shape `[seq_len, batch_size, d_model]`
      """
  
      # Prepare `query`, `key` and `value` for attention computation
      # `key` and `value`  will then have shape `[seq_len, batch_size, heads, d_k]`
      # and `query` will have shape `[batch_size, heads, d_k]`
      query = self.query(query)
      if self.key:
          key = self.key(key)
      if self.value:
          value = self.value(value)
  
      # Compute attention scores.
      # Results in a tensor of shape `[seq_len, batch_size, heads]`
      scores = self.get_scores(query, key)
  
      # Scale scores $\frac{1}{\sqrt{d_k}}$
      scores *= self.scale
  
      # Softmax
      attn = self.softmax(scores)
  
      # Apply dropout
      attn = self.dropout(attn)
  
      # Multiply by the values
      x = torch.einsum("jbh,jbhd->bhd", attn, value)
  
      # Concatenate multiple heads
      x = x.reshape(x.shape[0], -1)
  
      # Output layer
      return self.output(x)
  
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 3 for calling the FeedbackAttention obj's underlying member method code with example usage: obj.forward() and returns: self