code: def _create_dataset_split(
      cfg: AttrDict,
      data_split: str,
      features_dim: int,
      kmeans,
      pca: Optional[PCA] = None,
  ):
      """
      Scan the dataset split and create a new classification dataset out of it
      where each image is associated to the centroid the closest in feature space.
      """
      num_clusters = cfg.CLUSTERFIT.NUM_CLUSTERS
      data_name = cfg.CLUSTERFIT.FEATURES.DATASET_NAME
      layer_name = cfg.CLUSTERFIT.FEATURES.LAYER_NAME
  
      logging.info(
          f"Computing cluster label assignment for each sample in {data_split}..."
      )
      indices = []
      distances = []
      target_clusters = []
      shard_paths = ExtractedFeaturesLoader.get_shard_file_names(
          input_dir=cfg.CLUSTERFIT.FEATURES.PATH,
          split=data_split.lower(),
          layer=cfg.CLUSTERFIT.FEATURES.LAYER_NAME,
      )
      for shard_path in shard_paths:
          shard_content = ExtractedFeaturesLoader.load_feature_shard(
              shard_path
          )
          shard_features = shard_content.features
  
          # TODO - factorize this with above??? normalization at least???
          # Reshape and normalize the loaded features
          shard_features = shard_features.reshape(
              shard_features.shape[0], -1
          )
          shard_features_norm = (
              np.linalg.norm(shard_features, axis=1) + 1e-5
          )
          shard_features = (
              shard_features / shard_features_norm[:, np.newaxis]
          )
  
          if pca is not None:
              shard_features = pca.transform(shard_features)
              shard_features = np.ascontiguousarray(shard_features)
          shard_distances, shard_cluster_labels = kmeans.index.search(
              shard_features, 1
          )
          indices.extend(shard_content.indices)
          distances.extend(shard_distances)
          target_clusters.extend(shard_cluster_labels)
  
      # Step 5: save clustering data and hard cluster labels for the images
      logging.info(
          "Saving centroids and cluster assignments to file..."
      )
      dataset_image_paths = get_image_paths(cfg, split=data_split)
      image_paths = [dataset_image_paths[i] for i in indices]
      data_split = data_split.lower()
      clustering_output_dict = {
          "sample_indices": indices,
          "hard_labels": target_clusters,
          "centroids": kmeans.centroids,
          "distances": distances,
          "images": image_paths,
      }
      output_dir = cfg.CLUSTERFIT.OUTPUT_DIR
      g_pathmgr.mkdirs(output_dir)
      output_prefix = f"{data_name}_{data_split}_{layer_name}_N{num_clusters}_D{features_dim}"
      cluster_output_filepath = os.path.join(
          output_dir, f"{output_prefix}.pkl"
      )
      labels_output_filepath = os.path.join(
          output_dir, f"{output_prefix}_labels.npy"
      )
      image_path_filepath = os.path.join(
          output_dir, f"{output_prefix}_images.npy"
      )
      out_images = np.array(image_paths)
      out_hard_labels = np.array(
          target_clusters, dtype=np.int64
      ).reshape(-1)
      save_file(clustering_output_dict, cluster_output_filepath)
      save_file(out_images, image_path_filepath)
      save_file(out_hard_labels, labels_output_filepath)
  
imports: import logging
  import os
  import sys
  from argparse import Namespace
  from typing import Any, List, Optional
  import numpy
  import torch
  from iopath.common.file_io import g_pathmgr
  from vissl.config import AttrDict
  from vissl.data import build_dataset
  from vissl.hooks import default_hook_generator
  from vissl.utils.checkpoint import get_checkpoint_folder
  from vissl.utils.distributed_launcher import launch_distributed
  from vissl.utils.env import set_env_vars, setup_path_manager
  from vissl.utils.extract_features_utils import ExtractedFeaturesLoader
  from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict
  from vissl.utils.io import save_file
  from vissl.utils.logger import setup_logging, shutdown_logging
  from vissl.utils.misc import is_faiss_available, set_seeds
  from vissl.utils.pca import PCA
label_desc: Use image node_id 4 for calling a global function with example usage: _create_dataset_split(cfg, data_split, features_dim, kmeans, pca) without return types