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","."]
docstr: implementation of :class:`~maze.perception.blocks.base.PerceptionBlock` interface
docstr_tok: ["implementation","of",":","class",":","`","~maze.perception.blocks.base.PerceptionBlock","`","interface"]
code: def forward(
      self, block_input: Dict[str, torch.Tensor]
  ) -> Dict[str, torch.Tensor]:
      """implementation of :class:`~maze.perception.blocks.base.PerceptionBlock` interface"""
  
      # input_tensor.shape = (B X C X W X H)
      input_tensor = block_input[self.in_keys[0]]
      attn_mask = (
          block_input[self.in_keys[1]]
          if len(self.in_keys) > 1
          else None
      )
  
      assert len(input_tensor.shape) == 4
      batch_size, num_channels, width, height = input_tensor.size()
  
      # N = width * height, E = embed_dim
      proj_query = (
          self.query_conv(input_tensor)
          .view(batch_size, -1, width * height)
          .permute(0, 2, 1)
      )  # B X E X (N)
      proj_key = self.key_conv(input_tensor).view(
          batch_size, -1, width * height
      )  # B X E x (N)
      energy = torch.bmm(proj_query, proj_key)  # transpose check
  
      if attn_mask is not None:
          energy += attn_mask
  
      attention = self.softmax(energy)  # B X (N) X (N)
      attention = self.dropout(attention)
  
      proj_value = self.value_conv(input_tensor).view(
          batch_size, -1, width * height
      )  # B X C X N
  
      out = torch.bmm(proj_value, attention.permute(0, 2, 1))
      out = out.view(batch_size, num_channels, width, height)
  
      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 seconds 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_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 2 for calling the SelfAttentionConvBlock obj's underlying member method code with example usage: obj.forward(block_input) and returns: out_dict