class_docstr: Gets knn of test features from all processes on a chunk of the train features
  
  Each rank gets a chunk of the train features as well as a chunk of the test features.
  In `compute_neighbors`, for each rank one after the other, its chunk of test features
  is sent to all devices, partial knns are computed with each chunk of train features
  then collated back on the original device.
code: def _gather_all_knn_for_rank(
      self, topk_sims, neighbors_labels, target_rank
  ):
      # Gather all neighbors for `target_rank`
      topk_sims_rank = retrieved_rank = None
      if self.global_rank == target_rank:
          topk_sims_rank = [
              torch.zeros_like(topk_sims)
              for _ in range(self.global_size)
          ]
          retrieved_rank = [
              torch.zeros_like(neighbors_labels)
              for _ in range(self.global_size)
          ]
  
      torch.distributed.gather(
          topk_sims, topk_sims_rank, dst=target_rank
      )
      torch.distributed.gather(
          neighbors_labels, retrieved_rank, dst=target_rank
      )
  
      if self.global_rank == target_rank:
          # Perform a second top-k on the k * global_size retrieved neighbors
          topk_sims_rank = torch.cat(topk_sims_rank, dim=1)
          retrieved_rank = torch.cat(retrieved_rank, dim=1)
          results = self._get_knn_sims_and_labels(
              topk_sims_rank, retrieved_rank
          )
          return results
      return None
  
imports: import argparse
  from functools import partial
  import json
  import logging
  import os
  import sys
  from typing import List, Optional
  import torch
  from torch.nn.functional import one_hot, softmax
  import dinov2.distributed
  from dinov2.data import SamplerType, make_data_loader, make_dataset
  from dinov2.data.transforms import make_classification_eval_transform
  from dinov2.eval.metrics import AccuracyAveraging, build_topk_accuracy_metric
  from dinov2.eval.setup import get_args_parser
  from dinov2.eval.setup import setup_and_build_model
  from dinov2.eval.utils import ModelWithNormalize, evaluate, extract_features
label_desc: Use image node_id 4 for calling the KnnModule obj's underlying member method code with example usage: obj._gather_all_knn_for_rank(topk_sims, neighbors_labels, target_rank) and returns: None, results