code: def __init__(
      self,
      *,
      d_model: int,
      attn: FastWeightsAttention,
      feed_forward: FeedForward,
      dropout_prob: float,
  ):
      super().__init__()
      # Transformer size $d_{model}$
      self.size = d_model
      #
      self.attn = attn
      self.feed_forward = feed_forward
      self.dropout = nn.Dropout(dropout_prob)
  
      # Normalization layers
      self.norm_self_attn = nn.LayerNorm([d_model])
      self.norm_ff = nn.LayerNorm([d_model])
  
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 FastWeightsAttentionTransformerLayer object from inherited base classes: Module with example: obj = FastWeightsAttentionTransformerLayer()