class_docstr: Implementation of a self-attention block as described by reference: https://arxiv.org/abs/1805.08318
  
  This block can then be used for 2d data (images), to compute the self attention. If two out_keys are given, the
  actual attention is returned from the forward pass with the second out_key. Otherwise only the computed
  self-attention is returned
  
  :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 embed_dim: The embedding dimensionality, which should be an even fraction of the input channels.
  :param add_input_to_output: Specifies weather the computed self attention is added to the input and returned.
  :param bias: Specify weather to use a bias in the projections.
class_docstr_tok: ["Implementation","of","a","self-attention","block","as","described","by","reference",":","https",":","\/\/arxiv.org\/abs\/1805.08318","This","block","can","then","be","used","for","2d","data","(","images",")",",","to","compute","the","self","attention",".","If","two","out_keys","are","given",",","the","actual","attention","is","returned","from","the","forward","pass","with","the","second","out_key",".","Otherwise","only","the","computed","self-attention","is","returned",":","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","embed_dim",":","The","embedding","dimensionality",",","which","should","be","an","even","fraction","of","the","input","channels",".",":","param","add_input_to_output",":","Specifies","weather","the","computed","self","attention","is","added","to","the","input","and","returned",".",":","param","bias",":","Specify","weather","to","use","a","bias","in","the","projections","."]
code: def __init__(
      self,
      in_keys: Union[str, List[str]],
      out_keys: Union[str, List[str]],
      in_shapes: Union[Sequence[int], List[Sequence[int]]],
      embed_dim: int,
      dropout: Optional[float],
      add_input_to_output: bool,
      bias: bool,
  ):
      super().__init__(
          in_keys=in_keys, out_keys=out_keys, in_shapes=in_shapes
      )
  
      # Assertions
      assert len(self.in_keys) == len(self.in_shapes)
      assert len(self.in_keys) in (1, 2)
      # If two out keys are given, the seconds one is the attention
      assert len(self.out_keys) in (1, 2)
      assert (
          len(self.in_shapes[0]) == 3
      ), "In dimensionality should be 3 without batch"
      if len(self.in_keys) > 1:
          assert len(self.in_shapes[1]) == 2
          assert (
              np.prod(self.in_shapes[0][-2:])
              == self.in_shapes[1][-1]
              == self.in_shapes[1][-2]
          )
      in_dim = self.in_shapes[0][0]
      assert in_dim > embed_dim
      assert in_dim // embed_dim == in_dim / float(
          embed_dim
      ), "in_dim should be evenly dividable by embed_dim"
  
      self.in_dim = in_dim
      self.embedding_dim = embed_dim
      self.add_input_to_output = add_input_to_output
  
      self.query_conv = nn.Conv2d(
          in_channels=self.in_dim,
          out_channels=self.embedding_dim,
          kernel_size=1,
          bias=bias,
      )
      self.key_conv = nn.Conv2d(
          in_channels=self.in_dim,
          out_channels=self.embedding_dim,
          kernel_size=1,
          bias=bias,
      )
      self.value_conv = nn.Conv2d(
          in_channels=self.in_dim,
          out_channels=self.in_dim,
          kernel_size=1,
          bias=bias,
      )
      self.gamma = nn.Parameter(torch.zeros(1, dtype=torch.float32))
      self.softmax = nn.Softmax(dim=-1)
      self.dropout = nn.Dropout(
          p=dropout if dropout is not None else 0.0
      )
  
file_path: maze/maze/perception/blocks/general/self_attention_conv.py
imports: from typing import Union, List, Sequence, Dict, Optional
  import numpy
  import torch
  from torch import nn
  from maze.core.annotations import override
  from maze.perception.blocks.base import PerceptionBlock
label_desc: Use image node_id 1 to create a new SelfAttentionConvBlock object from inherited base classes: PerceptionBlock with example: obj = SelfAttentionConvBlock(in_keys, out_keys, in_shapes, embed_dim, dropout, add_input_to_output, bias)