class_docstr: A block applying global pooling with optional masking.
  Pooling is applied w.r.t. the selected pooling function, the pooling dimension and the optional mask (in_keys[1]).
  
  :param in_keys: One key identifying the input tensor, and an optional second one identifying the mask tensor.
  :param out_keys: One key identifying the output tensors.
  :param in_shapes: List of input shapes.
  :param pooling_func: Options: {'mean', 'sum', 'max'}.
  :param pooling_dim: The dimension(s) along which the pooling functions get applied.
class_docstr_tok: ["A","block","applying","global","pooling","with","optional","masking",".","Pooling","is","applied","w.r.t",".","the","selected","pooling","function",",","the","pooling","dimension","and","the","optional","mask","(","in_keys","[","1","]",")",".",":","param","in_keys",":","One","key","identifying","the","input","tensor",",","and","an","optional","second","one","identifying","the","mask","tensor",".",":","param","out_keys",":","One","key","identifying","the","output","tensors",".",":","param","in_shapes",":","List","of","input","shapes",".",":","param","pooling_func",":","Options",":","{","'mean","'",",","'sum","'",",","'max","'","}",".",":","param","pooling_dim",":","The","dimension","(","s",")","along","which","the","pooling","functions","get","applied","."]
docstr: Forward pass through the block, iterating over the first 2 dimensions and pooling the rest in dim=0.
  
  :param block_input: The block's input dictionary.
  :return: The block's output dictionary.
docstr_tok: ["Forward","pass","through","the","block",",","iterating","over","the","first","2","dimensions","and","pooling","the","rest","in","dim=0",".",":","param","block_input",":","The","block","'s","input","dictionary",".",":","return",":","The","block","'s","output","dictionary","."]
code: def forward(
      self, block_input: Dict[str, torch.Tensor]
  ) -> Dict[str, torch.Tensor]:
      """Forward pass through the block, iterating over the first 2 dimensions and pooling the rest in dim=0.
  
      :param block_input: The block's input dictionary.
      :return: The block's output dictionary.
      """
  
      # check input tensor
      input_tensor = block_input[self.in_keys[0]]
      if self._use_masking:
          mask_tensor = block_input[self.in_keys[1]]
  
          # prepare mask tensor
          mask_tensor: torch.Tensor = ~torch.eq(mask_tensor, 0)
      else:
          mask_tensor = torch.ones(1).to(input_tensor.device)
  
      # apply pooling
      rep = self.pooling_func(
          input_tensor=input_tensor,
          mask_tensor=mask_tensor,
          dim=self.pooling_dim,
      )
  
      return {self.out_keys[0]: rep}
  
file_path: maze/maze/perception/blocks/general/masked_global_pooling.py
imports: from typing import Union, List, Sequence, Dict
  import torch
  from maze.core.annotations import override
  from maze.perception.blocks.base import PerceptionBlock
label_desc: Use image node_id 5 for calling the MaskedGlobalPoolingBlock obj's underlying member method code with example usage: obj.forward(block_input) and returns: dict