class_docstr: ## Updated Feedback Transformer Module
  
  This is the updated feedback transformer module that caches the keys and values.
class_docstr_tok: ["#","#","Updated","Feedback","Transformer","Module","This","is","the","updated","feedback","transformer","module","that","caches","the","keys","and","values","."]
docstr: * `layer` is the feedback transformer layer, which we clone for each layer
  * `n_layers` is the number of layers in the transformer
  * `d_model` is the number of features in the transformer
  * 'heads' is the number of attention heads
docstr_tok: ["*","`","layer","`","is","the","feedback","transformer","layer",",","which","we","clone","for","each","layer","*","`","n_layers","`","is","the","number","of","layers","in","the","transformer","*","`","d_model","`","is","the","number","of","features","in","the","transformer","*","'heads","'","is","the","number","of","attention","heads"]
code: def __init__(
      self,
      layer: FeedbackTransformerLayer,
      n_layers: int,
      d_model: int,
      heads: int,
  ):
      """
      * `layer` is the feedback transformer layer, which we clone for each layer
      * `n_layers` is the number of layers in the transformer
      * `d_model` is the number of features in the transformer
      * 'heads' is the number of attention heads
      """
  
      super().__init__()
      # Make copies of the transformer layer
      self.layers = clone_module_list(layer, n_layers)
      # Final normalization layer
      self.norm = nn.LayerNorm([layer.size])
      # Memory vectors are computed as a weighted sum of representations of each layer.
      # This is the weights parameter for that.
      self.weights = nn.Parameter(
          torch.ones(n_layers + 1), requires_grad=True
      )
      # Softmax for weights before taking the weighted sum
      self.softmax = nn.Softmax(0)
  
      # Number of features in a head
      d_k = d_model // heads
      # Module to transform embeddings (memory) to get keys
      self.key = PrepareForMultiHeadAttention(
          d_model, heads, d_k, bias=False
      )
      # Module to transform embeddings (memory) to get keys
      self.value = PrepareForMultiHeadAttention(
          d_model, heads, d_k, bias=False
      )
  
      # Memory for stacked keys
      self.mem_key = Stack(512)
      # Memory for stacked values
      self.mem_value = Stack(512)
  
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 FeedbackTransformerKV object from inherited base classes: Module with example: obj = FeedbackTransformerKV(layer, n_layers, d_model, heads)