class_docstr: Implementation of a torch MultiHeadAttention block.
  
  This Block wraps the torch.nn.MultiheadAttention. This block can then be used
  for 1d data as well as sequential data.
  
  :param in_keys: Keys identifying the input tensors. First key is the query, second is the key and the third input
      is the value. Additionally there is the optional attention mask which can be passed as an input to the block.
      - query: :math:`(N, L, E)` where L is the target sequence length, N is the batch size, E is
      the embedding dimension.
      - key: :math:`(N, S, E)`, where S is the source sequence length, N is the batch size, E is
      the embedding dimension.
      - value: :math:`(N, S, E)` where S is the source sequence length, N is the batch size, E is
      the embedding dimension.
      - (Optional) attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence
      length. 3D mask :math:`(N, L, S)` where N is the batch size, L is the target sequence length, S is the
      source sequence length. attn_mask ensure that position i is allowed to attend the unmasked positions. If a
      Bool- or int- or float-Tensor is provided, positions with ``False/~1`` is not allowed to attend while
      ``True/1`` values will be unchanged.
      - (Optional) key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
      If a Bool- or int- or float-Tensor is provided, the positions with the value of ``False/~1`` will be ignored
      while the position with the value of ``True/1`` will be unchanged.
  :param out_keys: Keys identifying the output tensors. First key is self-attention output, second optional key is
      attention map.
      - attn_output: :math:`(N, L, E)` where L is the target sequence length, N is the batch size, E is the embedding
      dimension.
      - attn_output_weights: :math:`(N, L, S)` where N is the batch size, L is the target sequence length, S is the
      source sequence length.
  :param in_shapes: List of input shapes.
  :param num_heads: Parallel attention heads.
  :param dropout: A dropout layer on attn_output_weights.
  :param bias: Add bias as module parameter.
  :param add_bias_kv: Add bias to the key and value sequences at dim=0.
  :param add_zero_attn: Add a new batch of zeros to the key and value sequences at dim=1.
  :param kdim: Total number of features in key. Default: None.
  :param vdim: Total number of features in value. Default: None.
  :param use_key_padding_mask: Specify wether a key padding mask is being used.
  
  Note: If kdim and vdim are None, they will be set to embed_dim such that query, key, and value have the same number
      of features.
code: def __repr__(self):
      txt = f"{self.__class__.__name__}"
      txt += f"\n\tnum_heads: {self.self_attn.num_heads}"
      txt += f"\n\tembed_dim: {self.self_attn.embed_dim}"
      txt += f"\n\tdropout: {self.self_attn.dropout}"
      txt += f"\n\tbias: {self.self_attn.in_proj_bias is not None}"
      txt += f"\n\tadd_input_to_output: {self.add_input_to_output}"
      txt += (
          f"\n\tuse_attn_mask: "
          f"{len(self.in_keys) == 4 and not self.use_key_padding_mask or len(self.in_keys) == 5}"
      )
      txt += (
          f"\n\tuse_key_padding_mask: "
          f"{len(self.in_keys) == 4 and self.use_key_padding_mask or len(self.in_keys) == 5}"
      )
      txt += f"\n\tadd_bias_kv: {self.add_bias_kv}"
      txt += f"\n\tadd_zero_attn: {self.add_zero_attn}"
      if self.kdim is not None:
          txt += f"\n\tkdim: {self.kdim}"
      if self.vdim is not None:
          txt += f"\n\tvdim: {self.vdim}"
      txt += f"\n\tOut Shapes: {self.out_shapes()}"
      return txt
  
file_path: maze/maze/perception/blocks/feed_forward/multi_head_attention.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 3 for calling the MultiHeadAttentionBlock obj's underlying member method code with example usage: obj.__repr__() and returns: txt