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: #### Create local mask
  
  This creates a mask for
  
  egin{align}
  m_{t,t'} =
  egin{cases}
  1,  & { ext{for } \lvert t-t' 
  vert \lt s} \
  0, &    ext{otherwise}
  \end{cases}
  \end{align}
code: def create_local_mask(seq_len, local_window_size):
      """
          #### Create local mask
  
          This creates a mask for
  
          \begin{align}
          m_{t,t'} =
          \begin{cases}
          1,  & {\text{for } \lvert t-t' \rvert \lt s} \\
          0, & \text{otherwise}
          \end{cases}
          \end{align}
          """
  
      # Initialize to ones
      local_mask = torch.ones(seq_len, seq_len, dtype=torch.bool)
      # Make $t' - t \ge s$ zero
      local_mask = torch.tril(local_mask, local_window_size - 1)
      # Make $t - t' \ge s$ zero
      local_mask = torch.triu(local_mask, -(local_window_size - 1))
  
      #
      return local_mask
  
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 2 for calling the AFTLocal obj's underlying member method code with example usage: obj.create_local_mask(seq_len, local_window_size) and returns: local_mask