code: def kmeans(
      k,
      data,
      niter=25,
      seed=1234,
      checkpoint=None,
      verbose=True,
      return_stats=False,
  ):
      """Pure python kmeans implementation. Follows the Faiss C++ version
      quite closely, but takes a DatasetAssign instead of a training data
      matrix. Also redo is not implemented."""
      n, d = data.count(), data.dim()
  
      log = print if verbose else print_nop
  
      log(
          (
              "Clustering %d points in %dD to %d clusters, "
              + "%d iterations seed %d"
          )
          % (n, d, k, niter, seed)
      )
  
      rs = np.random.RandomState(seed)
      print("preproc...")
      t0 = time.time()
      # initialization
      perm = rs.choice(n, size=k, replace=False)
      centroids = data.get_subset(perm)
  
      iteration_stats = []
  
      log("  done")
      t_search_tot = 0
      obj = []
      for i in range(niter):
          t0s = time.time()
  
          log("assigning", end="\r", flush=True)
          assign, D, sums = data.assign_to(centroids)
  
          log("compute centroids", end="\r", flush=True)
  
          t_search_tot += time.time() - t0s
  
          err = D.sum()
          obj.append(err)
  
          hassign = np.bincount(assign, minlength=k)
  
          fac = hassign.reshape(-1, 1).astype("float32")
          fac[fac == 0] = 1  # quiet warning
  
          centroids = sums / fac
  
          nsplit = reassign_centroids(hassign, centroids, rs)
  
          s = {
              "obj": err,
              "time": (time.time() - t0),
              "time_search": t_search_tot,
              "imbalance_factor": imbalance_factor(k, assign),
              "nsplit": nsplit,
          }
  
          log(
              (
                  "  Iteration %d (%.2f s, search %.2f s): "
                  "objective=%g imbalance=%.3f nsplit=%d"
              )
              % (
                  i,
                  s["time"],
                  s["time_search"],
                  err,
                  s["imbalance_factor"],
                  nsplit,
              )
          )
          iteration_stats.append(s)
  
          if checkpoint is not None:
              log("storing centroids in", checkpoint)
              np.save(checkpoint, centroids)
  
      if return_stats:
          return centroids, iteration_stats
      else:
          return centroids
  
imports: import numpy
  import faiss
  import time
  from multiprocessing.pool import ThreadPool
label_desc: Use image node_id 8 for calling a global function with example usage: kmeans(k, data, niter, seed, checkpoint, verbose, return_stats) and returns: centroids,  iteration_stats, centroids