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 __init__(
      self, heads: int, d_model: int, dropout_prob: float, phi: DPFP
  ):
      super().__init__()
  
      # Number of features per head $d_k$
      self.d_k = d_model // heads
      # Number of heads
      self.heads = heads
  
      # These transform the `query`, `key` and `value` multi-headed attention.
      self.query = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=False
      )
      self.key = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=False
      )
      self.value = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=False
      )
  
      # Interpolation weight function $\sigma \Big(\textcolor{orange}{W_\beta} x^{(i)} \Big)$ for each head
      self.interpolation_weight = nn.Sequential(
          PrepareForMultiHeadAttention(d_model, heads, 1, bias=False),
          nn.Sigmoid(),
      )
  
      # $\textcolor{lightgreen}{\phi'}$
      self.phi = phi
  
      # Output layer
      self.output = nn.Linear(d_model, d_model)
      # Dropout
      self.dropout = nn.Dropout(dropout_prob)
  
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 1 to create a new FastWeightsAttention object from inherited base classes: Module with example: obj = FastWeightsAttention(heads, d_model, dropout_prob, phi)