code: def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor]):
      query = self.phi(self.query(x))
      key = self.phi(self.key(x))
      value = self.value(x)
  
      if weights is None:
          weights = key.new_zeros(
              (key.shape[0], key.shape[1], value.shape[2], key.shape[2])
          )
  
      value_existing = torch.einsum("bhvk,bhk->bhv", weights, key)
  
      beta = self.gate(x)
  
      weights = weights + torch.einsum(
          "bhv,bhk->bhvk", beta * (value - value_existing), key
      )
  
      x = torch.einsum("bhvk,bhk->bhv", weights, query)
  
      # Concatenate multiple heads
      x = x.reshape(x.shape[0], -1)
  
      # Output layer
      return self.output(x), weights
  
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 2 for calling the FastWeightsAttention obj's underlying member method code with example usage: obj.forward(x, weights) and returns: self,  weights