code: def balanced_sub_sampling(
      labels: np.ndarray,
      num_samples: int,
      skip_samples: int = 0,
      seed: int = 0,
  ) -> np.ndarray:
      """
      Given all the labels of a dataset, sub_sample a part of the labels such that:
      - the number of samples of each label differs by at most one
      - the sub sampling is deterministic (identical across distributed workers)
      Return the indices of the selected labels
      """
      groups = {}
      for i, label in enumerate(labels):
          groups.setdefault(label, []).append(i)
  
      unique_labels = sorted(groups.keys())
      skip_quotient, skip_rest = divmod(
          skip_samples, len(unique_labels)
      )
      sample_quotient, sample_rest = divmod(
          num_samples, len(unique_labels)
      )
      assert (
          sample_quotient > 0
      ), "the number of samples should be at least equal to the number of classes"
  
      with with_temporary_numpy_seed(seed):
          for i, label in enumerate(unique_labels):
              label_indices = groups[label]
              num_label_samples = sample_quotient + (
                  1 if i < sample_rest else 0
              )
              skip_label_samples = skip_quotient + (
                  1 if i < skip_rest else 0
              )
              permuted_indices = np.random.choice(
                  label_indices,
                  size=skip_label_samples + num_label_samples,
                  replace=False,
              )
              groups[label] = permuted_indices[skip_label_samples:]
  
      return np.concatenate([groups[label] for label in unique_labels])
  
imports: import contextlib
  import logging
  import queue
  from typing import Dict, List
  import numpy
  import torch
  from PIL import Image
  from torch.utils.data import Dataset
  from torch.utils.data.distributed import DistributedSampler
label_desc: Use image node_id 4 for calling a global function with example usage: balanced_sub_sampling(labels, num_samples, skip_samples, seed) and returns: np