class_docstr: A repeat-to-match block.
  This blocks takes two tensors and a dimension index as an input. Then when it's
  forward method is called, it matches the specified the dimension (with :param repeat_at_idx) of the first
  tensor with the specified dimension (:param repeat_at_idx)  of the second tensor. This is done by repeating the
  first tensor n times in dimension in dimension :param repeat_at_idx.
  Here n = tensor_1.shape[repeat_at_idx] - tensor_0.shape[repeat_at_idx]. As a constraint the first tensor has to
  satisfy the following condition: tensor_0[repeat_at_idx] == 1
  
  :param in_keys: The keys identifying the input tensors.
  :param out_keys: One key identifying the output tensors.
  :param in_shapes: List of input shapes.
  :param repeat_at_idx: Specify the dimension that should be matched between the tensors.
docstr: Forward pass, repeat the first tensor to match the second one in the given dimension
  
  :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, repeat the first tensor to match the second one in the given dimension
  
      :param block_input: The block's input dictionary.
      :return: The block's output dictionary.
      """
  
      # check input tensor
      tensor_to_repeat = block_input[self.in_keys[0]]
      assert tensor_to_repeat.shape[self.repeat_at_idx] == 1, (
          f"tensor_0.shape[self.repat_at_idx] should be 1, but "
          f"got {tensor_to_repeat.shape[self.repeat_at_idx]} ("
          f"full shape: {tensor_to_repeat.shape})"
      )
  
      num_of_repeats = block_input[self.in_keys[1]].shape[
          self.repeat_at_idx
      ]
  
      # forward pass
      repeats = list([1 for _ in range(tensor_to_repeat.ndim)])
      repeats[self.repeat_at_idx] = num_of_repeats
  
      output_tensor = tensor_to_repeat.repeat(*repeats)
  
      return {self.out_keys[0]: output_tensor}
  
file_path: maze/maze/perception/blocks/general/repeat_to_match.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 2 for calling the RepeatToMatchBlock obj's underlying member method code with example usage: obj.forward(block_input) and returns: dict