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: * `x_seq` is the input with shape `[seq_len, batch_size, d_model]`
docstr_tok: ["*","`","x_seq","`","is","the","input","with","shape","`","[","seq_len",",","batch_size",",","d_model","]","`"]
code: def forward(self, x_seq: torch.Tensor):
      """
      * `x_seq` is the input with shape `[seq_len, batch_size, d_model]`
      """
  
      # Split the input to a list along the sequence axis
      x_seq = torch.unbind(x_seq, dim=0)
      # List to store the outputs
      res = []
      # For each input step
      for step, x in enumerate(x_seq):
          # List to store layer outputs
          layer_outputs = [x]
  
          # Stack of keys and values
          key_tensor = None
          value_tensor = None
          # Get the keys and values tensors if we are beyond the initial step
          if step > 0:
              key_tensor = self.mem_key.get()
              value_tensor = self.mem_value.get()
  
          # Run through each layer
          for layer in self.layers:
              # Get layer output
              x = layer(x=x, key=key_tensor, value=value_tensor)
              # Append them to the list of layer outputs
              layer_outputs.append(x)
  
          # Stack the layer outputs to a tensor
          layer_outputs = torch.stack(layer_outputs)
          # Calculate the memory vector as a weighted sum of layer outputs
          mem = torch.einsum(
              "lbd,l->bd", layer_outputs, self.softmax(self.weights)
          )
          # Calculate the keys from memory and add it to the stack
          self.mem_key.append(step, self.key(mem))
          # Calculate the values from memory and add it to the stack
          self.mem_value.append(step, self.value(mem))
          # Append the output to results
          res.append(x)
  
      # Stack the output tensors
      res = torch.stack(res)
      # Normalize the output
      return self.norm(res)
  
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 FeedbackTransformerKV obj's underlying member method code with example usage: obj.forward(x_seq) and returns: self