code: def get_closest_embeddings(
      affinity_mat: torch.Tensor, n_closest: int
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Get the indices of the embedding vectors we want to merge.
  
      Example:
          >>> n_closest = 2
          >>> affinity_mat = [[1.0, 0.2, 0.8],
                              [0.2, 1.0, 0.4],
                              [0.8, 0.4, 1.0]]
          >>> affinity_mat.sum(0)
          [2.0, 1.6, 2.2]
  
          # The closest two embedding vectors are at index 0 and 2.
  
      Args:
          affinity_mat: (Tensor)
              Symmetric affinity matrix of the given embedding vector set.
          n_closest (int):
              The amount of vector counts that are expected to be removed from the set
              Example:
                  Input: 10 vectors in a set
                  n_closest = 5
                  (5+1) vectors are merged into 1 vector
                  Output: 5 vectors in a set
  
      Returns:
          idx_aff_sum (torch.Tensor):
              Indices of the closest `n_closest` embedding vectors
          rest_inds (torch.Tensor):
              Indices of the complementary set of the indices in `idx_aff_sum`
      """
      comb_limit = int(affinity_mat.shape[0] - 1)
      if n_closest > comb_limit:
          raise ValueError(
              f"Got n_closest of {n_closest}: {n_closest} is bigger than comb_limit {comb_limit}"
          )
  
      # Take summed values over one axis
      sum_cmat = affinity_mat.sum(0)
  
      # `n_closest + 1` will become 1 embedding vector after merging
      idx_aff_sum = torch.argsort(sum_cmat, descending=True)[
          : (n_closest + 1)
      ]
      rest_inds = torch.argsort(sum_cmat, descending=True)[
          (n_closest + 1) :
      ]
      return idx_aff_sum, rest_inds
  
imports: from typing import List, Tuple
  import torch
  from nemo.collections.asr.parts.utils.offline_clustering import NMESC, SpeakerClustering, SpectralClustering, get_scale_interpolated_embs, getAffinityGraphMat, getCosAffinityMatrix, split_input_data
  from nemo.collections.asr.parts.utils.optimization_utils import linear_sum_assignment
label_desc: Use image node_id 7 for calling a global function with example usage: get_closest_embeddings(affinity_mat, n_closest) and returns: idx_aff_sum,  rest_inds