class_docstr: ### AFT Local Operation
  
  $$Y_t = \sigma(Q_t) \odot
   rac{\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'}) \odot V_{t'}}
   {\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'})}$$
  
  where,
  
  egin{align}
  w'_{t,t'} =
  egin{cases}
  w_{t,t'},  & {      ext{for } \lvert t-t' 
  vert \lt s} \
  0, &        ext{otherwise}
  \end{cases}
  \end{align}
docstr: `query`, `key` and `value` are the tensors that store
  collection of token embeddings for  *query*, *key* and *value*.
  They have shape `[seq_len, batch_size, d_model]`.
  
  `mask` has shape `[seq_len, seq_len, batch_size]` and
  `mask[i, j, b]` indicates whether for batch `b`,
  query at position `i` has access to key-value at position `j`.
code: def forward(
      self,
      *,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
  ):
      """
      `query`, `key` and `value` are the tensors that store
      collection of token embeddings for  *query*, *key* and *value*.
      They have shape `[seq_len, batch_size, d_model]`.
  
      `mask` has shape `[seq_len, seq_len, batch_size]` and
      `mask[i, j, b]` indicates whether for batch `b`,
      query at position `i` has access to key-value at position `j`.
      """
  
      # `query`, `key` and `value`  have shape `[seq_len, batch_size, d_model]`
      seq_len, _, _ = query.shape
  
      if mask is not None:
          # `mask` has shape `[seq_len_q, seq_len_k, batch_size]`,
          # where first dimension is the query dimension.
          # If the query dimension is equal to $1$ it will be broadcasted.
          assert mask.shape[0] == 1 or mask.shape[0] == query.shape[0]
          assert mask.shape[1] == key.shape[0]
          assert mask.shape[2] == 1 or mask.shape[2] == query.shape[1]
  
      # Transform query, key and value embeddings
      query = self.query(query)
      key = self.key(key)
      value = self.value(value)
  
      # Get
      #
      #     \begin{align}
      #     w'_{t,t'} =
      #     \begin{cases}
      #     w_{t,t'},  & {\text{for }\lvert t-t' \rvert \lt s} \\
      #     0, & \text{otherwise}
      #     \end{cases}
      #     \end{align}
      #
      # using the mask
      pos_bias = (
          self.pos_bias[:seq_len, :seq_len]
          * self.local_mask[:seq_len, :seq_len]
      )
      pos_bias = pos_bias.unsqueeze(-1)
      pos_bias.masked_fill_(~mask, float("-inf"))
  
      # \begin{align}
      # Y_t &= \sigma(Q_t) \odot
      # \frac{\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'}) \odot V_{t'}}
      # {\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'})} \\
      # &= \sigma(Q_t) \odot
      #    \frac{\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'}) \odot V_{t'}}
      #    {\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'})}
      # \end{align}
      #
      # We compute $\exp(w_{t,t'})$, $\exp(K_{t'}) \odot V_{t'}$ and $\exp(K_{t'})$
      # separately and do a matrix multiplication. We use einsum for clarity.
  
      # We subtract $\max_{t'}(K_{t'})$ and $\max_{t'}(w_{t,t'})$ before calculating the exponents to stabilize
      # the softmax calculation.
      #
      # If $x_i$ is large $\exp(x_i)$ becomes huge and the computation of
      # $\frac{\sum\exp(x_i)y_i}{\sum\exp(x_i)}$becomes unstable.
      # Subtracting a constant before calculating the exponent from numerator and denominator will cancel out.
      # and can help stabilize the computation.
      # So we subtract $\max(x_i)$ to stabilize the computation.
      max_key = key.max(dim=0, keepdims=True)[0]
      max_pos_bias = pos_bias.max(dim=1, keepdims=True)[0]
  
      # $\exp \big(K_{t'}- \max_{t'}(K_{t'})\big)$
      exp_key = torch.exp(key - max_key)
      # $\exp \big(w_{t,t'} - \max_{t'}(w_{t,t'})\big)$
      exp_pos_bias = torch.exp(pos_bias - max_pos_bias)
  
      # The numerator part $\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'}) \odot V_{t'}$
      num = torch.einsum("ijb,jbd->ibd", exp_pos_bias, exp_key * value)
      # The denominator part $\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'})$
      den = torch.einsum("ijb,jbd->ibd", exp_pos_bias, exp_key)
  
      # Output $$Y_t = \sigma(Q_t) \odot
      #         \frac{\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'}) \odot V_{t'}}
      #         {\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'})}$$
      y = self.activation(query) * num / den
  
      # Output layer
      return self.output(y)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/aft/__init__.py
imports: from typing import Optional
  import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 3 for calling the AFTLocal obj's underlying member method code with example usage: obj.forward() and returns: self