class_docstr: Implementation of a self-attention block as described by reference: https://arxiv.org/abs/1706.03762
  
  Within this block the torch nn.MuliheadAttention is used to model the self attention. This block can then be used
  for 1d data as well as sequential data, where the embedding dimensionality has to be equal to the last dimension
  of the input.
  
  :param in_keys: Keys identifying the input tensors. First key is self_attention output, second optional key is
      attention mask.
  :param out_keys: Keys identifying the output tensors. First key is self-attention output, second optional key is
      attention map.
  :param in_shapes: List of input shapes.
  :param num_heads: Parallel attention heads.
  :param dropout: A dropout layer on attn_output_weights.
  :param add_input_to_output: Specifies weather the computed self attention is added to the input and returned.
  :param bias: Add bias as module parameter.
docstr: implementation of :class:`~maze.perception.blocks.base.PerceptionBlock` interface
docstr_tok: ["implementation","of",":","class",":","`","~maze.perception.blocks.base.PerceptionBlock","`","interface"]
code: def normalized_forward(
      self, block_input: Dict[str, torch.Tensor]
  ) -> Dict[str, torch.Tensor]:
      """implementation of :class:`~maze.perception.blocks.base.PerceptionBlock` interface"""
  
      # Get the tensors from the block input
      input_tensor = block_input[self.in_keys[0]]
      attn_mask = (
          block_input[self.in_keys[1]]
          if len(self.in_keys) > 1
          else None
      )
  
      # If the attention mask is used, it has to have the shape (num_heads * batch_size, target sequence length,
      #   source sequence length). Thus we have to repeat for the number of heads.
      if attn_mask is not None:
          if self.num_heads > 1:
              attn_mask = attn_mask.repeat(
                  [self.num_heads, *[1 for _ in attn_mask.shape[1:]]]
              )
          # Furthermore we have to invert the mask in order to work with the torch.nn.MultiheadAttention
          attn_mask = ~torch.eq(
              attn_mask, torch.tensor(1).to(attn_mask.device)
          )
          # Finally the first value of the mask is set to true in oder to circumvent nan values, while still ensuring
          #   fast processing of the block.
          attn_mask[..., 0] = False
  
      if self.preprocess is not None:
          input_tensor = self.preprocess(input_tensor)
  
      out, attention = self.self_attn(
          input_tensor,
          input_tensor,
          input_tensor,
          need_weights=len(self.out_keys) == 2,
          attn_mask=attn_mask,
      )
  
      if self.postprocess is not None:
          out = self.postprocess(out)
          input_tensor = self.postprocess(input_tensor)
  
      # Scale self attention
      out = self.gamma * out
  
      if self.add_input_to_output:
          out = out + input_tensor
  
      out_dict = dict({self.out_keys[0]: out})
      # If two out keys are given, the second one is the attention
      if len(self.out_keys) == 2:
          out_dict[self.out_keys[1]] = attention
      return out_dict
  
file_path: maze/maze/perception/blocks/general/self_attention_seq.py
imports: from typing import Union, List, Sequence, Dict, Optional
  import torch
  from torch import nn
  from maze.core.annotations import override
  from maze.perception.blocks.shape_normalization import ShapeNormalizationBlock
label_desc: Use image node_id 2 for calling the SelfAttentionSeqBlock obj's underlying member method code with example usage: obj.normalized_forward(block_input) and returns: out_dict