class_docstr: _EfficientKMeans is primarily used by palettization to perform a k-means operation. This class also has an in-house
  implementation of k-means, called ``kmeans_pp`` which runs entirely on GPU and is ~10x faster than sklearn's API.
docstr: In-house implementation of kmeans that runs entirely on GPU and is ~10x faster.
docstr_tok: ["In-house","implementation","of","kmeans","that","runs","entirely","on","GPU","and","is","~10x","faster","."]
code: def kmeans_pp(
      self,
      n_init: str,
      X: _torch.Tensor,
      random_state: int,
      offset: int = 0,
  ):
      """
      In-house implementation of kmeans that runs entirely on GPU and is ~10x faster.
      """
      assert (
          len(X) >= self.n_clusters
      ), f"Weight fewer points than the number of clusters: {len(X)} vs. {self.n_clusters}"
  
      S = X[offset:]
  
      self.inertia_ = None
  
      width = (len(S) - 1) // (random_state + 1)
  
      for i in range(n_init):
          idx = int(i / n_init * width)
          C = S[idx].unsqueeze(0)
  
          for j in range(len(C), self.n_clusters):
              min_error, labels = self.__class__.x_c_dist(S, C).min(
                  dim=-1
              )
  
              while True:
                  max_dist_idx = _torch.argmax(min_error)
                  assert (
                      min_error[max_dist_idx] >= 0
                  ), "Cannot find a next candidate"
  
                  candidate_C = S[max_dist_idx]
                  if candidate_C in set(C):
                      _dist[max_dist_idx] = -1
                  else:
                      C = _torch.vstack((C, candidate_C))
                      break
  
          if len(set(C)) != self.n_clusters:
              return self.kmeans_pp(n_init, X, random_state, offset + 1)
  
          min_error, labels = self.__class__.x_c_dist(X, C).min(dim=-1)
          cur_cost = min_error.sum()
  
          if self.inertia_ is None or self.inertia_ > cur_cost:
              self.inertia_ = cur_cost
              self.cluster_centers_ = C
              self.labels_ = labels
  
file_path: coremltools/coremltools/optimize/torch/palettization/_efficient_kmeans.py
imports: import numpy
  import torch
  import torch.distributed
label_desc: Use image node_id 4 for calling the _EfficientKMeans obj's underlying member method code with example usage: obj.kmeans_pp(n_init, X, random_state, offset) and returns: self