code: def __init__(
      self, heads: int, d_model: int, dropout_prob: float, phi: DPFP
  ):
      super().__init__()
  
      # Number of features per head
      self.d_k = d_model // heads
      #
      self.heads = heads
  
      # These transform the `query` multi-headed attention.
      self.query = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=False
      )
      # These transform the `key` and `value` for multi-headed attention.
      self.key = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=False
      )
      self.value = PrepareForMultiHeadAttention(
          d_model, heads, self.d_k, bias=False
      )
  
      self.gate = nn.Sequential(
          PrepareForMultiHeadAttention(d_model, heads, 1, bias=False),
          nn.Sigmoid(),
      )
  
      self.phi = phi
  
      # Output layer
      self.output = nn.Linear(d_model, d_model)
      # Dropout
      self.dropout = nn.Dropout(dropout_prob)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/fast_weights/token_wise.py
imports: from typing import Optional
  import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.transformers.fast_weights import DPFP
  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)