code: def nearest_neighbors(
      input_maps: torch.Tensor,
      candidate_maps: torch.Tensor,
      distances: torch.Tensor,
      num_matches: int,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """Finds the nearest neighbors of the maps in input_maps in candidate_maps.
  
      Args:
          input_maps:
              A tensor of maps for which to find nearest neighbors.
              It has size: [batch_size, input_map_size, feature_dimension]
          candidate_maps:
              A tensor of maps to search for nearest neighbors.
              It has size: [batch_size, candidate_map_size, feature_dimension]
          distances:
              A tensor of distances between the maps in input_maps and candidate_maps.
              It has size: [batch_size, input_map_size, candidate_map_size]
          num_matches:
              Number of nearest neighbors to return. If num_matches is None or -1,
              all the maps in candidate_maps are considered.
  
      Returns:
          A tuple of tensors, containing the nearest neighbors in input_maps and candidate_maps.
          They both have size: [batch_size, input_map_size, feature_dimension]
      """
  
      if (
          num_matches is None
          or num_matches == -1
          or num_matches > input_maps.size(1)
      ):
          num_matches = input_maps.size(1)
  
      # Find nearest neighbour of each input element in the candidate map
      topk_values, topk_indices = distances.topk(
          k=1, dim=2, largest=False
      )  # [bsz, input_map_size, 1]
      topk_values = topk_values.squeeze(-1)  # [bsz, input_map_size]
  
      # Select num_matches neighbors pairs having the lowest distance value.
      _, min_indices = topk_values.topk(
          k=num_matches, dim=1, largest=False
      )  # [bsz, num_matches]
  
      # Create the filtered input map with num_matches lowest distance values.
      feature_dimension = input_maps.shape[2]
      filtered_input_maps = torch.gather(
          input_maps,
          1,
          min_indices.unsqueeze(-1).expand(-1, -1, feature_dimension),
      )  # [bsz, num_matches, feature_dimension]
  
      # Create candidate maps in the same way as input maps, but using corrispondent candidate values
      selected_candidate_maps = torch.gather(
          candidate_maps,
          1,
          topk_indices.expand(-1, -1, feature_dimension),
      )  # [bsz, input_map_size, feature_dimension]
      filtered_candidate_maps = torch.gather(
          selected_candidate_maps,
          1,
          min_indices.unsqueeze(-1).expand(-1, -1, feature_dimension),
      )  # [bsz, num_matches, feature_dimension]
  
      return filtered_input_maps, filtered_candidate_maps
  
imports: import math
  import warnings
  from typing import Iterable, List, Optional, Tuple, Union
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from numpy.typing import NDArray
  from torch.nn import Module
  from torch.nn.modules.batchnorm import _BatchNorm
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 19 for calling a global function with example usage: nearest_neighbors(input_maps, candidate_maps, distances, num_matches) and returns: filtered_input_maps,  filtered_candidate_maps