class_docstr: ## Feedback Transformer Module
class_docstr_tok: ["#","#","Feedback","Transformer","Module"]
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 = []
      # List to store the memory vectors
      mem = []
      # For each input step
      for x in x_seq:
          # List to store layer outputs
          layer_outputs = [x]
  
          # If there is memory, stack them into a vector
          mem_tensor = torch.stack(mem) if mem else None
  
          # Run through each layer
          for layer in self.layers:
              # Get layer output
              x = layer(x=x, key=mem_tensor, value=mem_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.append(
              torch.einsum(
                  "lbd,l->bd", layer_outputs, self.softmax(self.weights)
              )
          )
          # 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 FeedbackTransformer obj's underlying member method code with example usage: obj.forward(x_seq) and returns: self