class_docstr: ## Fast Weights Attention
  
  The paper introduces a new update rule for calculating $    extcolor{cyan}{W^{(i)}}$.
  The model first retrieves the current value
  $ar{v}^{(i)}$ paired with the key $k^{(i)}$.
  Then stores a combination $v^{(i)}_{new}$
  of the retrieved value $ar{v}^{(i)}$ and the input $v^{(i)}$.
  
  egin{align}
  k^{(i)}, v^{(i)}, q^{(i)} &=
      extcolor{orange}{W_k} x^{(i)},  extcolor{orange}{W_v} x^{(i)},  extcolor{orange}{W_q} x^{(i)} \
  ar{v}^{(i)} &=     extcolor{cyan}{W^{(i-1)}}       extcolor{lightgreen}{\phi'(k^{(i)})} \
  eta^{(i)} &= \sigma \Big(  extcolor{orange}{W_eta} x^{(i)} \Big) \
  v^{(i)}_{new} &= eta^{(i)} v^{(i)} + \Big(1 - eta^{(i)} \Big) ar{v}^{(i)} \
      extcolor{cyan}{W^{(i)}}
   &=         extcolor{cyan}{W^{(i-1)}} + v^{(i)}_{new} \otimes       extcolor{lightgreen}{\phi'(k^{(i)})} \
   &=         extcolor{cyan}{W^{(i-1)}} +
   eta^{(i)} \Big( v^{(i)} - ar{v}^{(i)} \Big ) \otimes     extcolor{lightgreen}{\phi'(k^{(i)})} \
  y^{(i)} &=  extcolor{cyan}{W^{(i)}}         extcolor{lightgreen}{\phi'(q^{(i)})}
  \end{align}
  
  where $     extcolor{orange}{W_eta}$ is a trainable parameter and $\sigma$ is the sigmoid function.
  
  Note that we don't need the normalization term $z$ because $        extcolor{lightgreen}{\phi'}$ is normalized.
code: def forward(self, x: torch.Tensor):
      # Get the number of steps $L$
      seq_len = x.shape[0]
      # $\textcolor{lightgreen}{\phi'(q^{(i)})}$ for all steps and heads
      query = self.phi(self.query(x))
      # $\textcolor{lightgreen}{\phi'(k^{(i)})}$ for all steps and heads
      key = self.phi(self.key(x))
      # $v^{(i)}$ for all steps and heads
      value = self.value(x)
      # $\beta^{(i)}$ for all steps and heads
      beta = self.interpolation_weight(x)
  
      # $\textcolor{cyan}{W^{(0)}}$
      weights = key.new_zeros(
          (key.shape[1], key.shape[2], value.shape[3], key.shape[3])
      )
      # List to store outputs $y^{(i)}$
      outputs = []
  
      # Iterate through steps
      for i in range(seq_len):
          # $$\bar{v}^{(i)} = \textcolor{cyan}{W^{(i-1)}} \textcolor{lightgreen}{\phi'(k^{(i)})}$$
          value_existing = torch.einsum(
              "bhvk,bhk->bhv", weights, key[i]
          )
  
          # $$\textcolor{cyan}{W^{(i)}}
          #      = \textcolor{cyan}{W^{(i-1)}} +
          #      \beta^{(i)} \Big( v^{(i)} - \bar{v}^{(i)} \Big ) \otimes \textcolor{lightgreen}{\phi'(k^{(i)})}$$
          weights = weights + torch.einsum(
              "bhv,bhk->bhvk",
              beta[i] * (value[i] - value_existing),
              key[i],
          )
  
          # $$y^{(i)} = \textcolor{cyan}{W^{(i)}} \textcolor{lightgreen}{\phi'(q^{(i)})}$$
          y = torch.einsum("bhvk,bhk->bhv", weights, query[i])
  
          # Merge multiple heads and append to `outputs`
          outputs.append(y.reshape(y.shape[0], -1))
  
      # Stack outputs at each step into a single tensor
      x = torch.stack(outputs)
  
      # Output layer
      return self.output(x)
  
imports: 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 FastWeightsAttention obj's underlying member method code with example usage: obj.forward(x) and returns: self