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: Method to run kmeans operation.
docstr_tok: ["Method","to","run","kmeans","operation","."]
code: def fit(self, X: _torch.Tensor):
      """
      Method to run kmeans operation.
      """
      N = len(X)
      if isinstance(self.cluster_centers_, str):
          if "kmeans++" in self.cluster_centers_:
              if _dist.is_available() and _dist.is_initialized():
                  world_size = _dist.get_world_size()
                  rank = _dist.get_rank()
              else:
                  world_size = 1
                  rank = 0
  
              if "cpu" in self.cluster_centers_:
                  import sklearn.cluster
  
                  kmeans = sklearn.cluster.KMeans(
                      n_init=max(10, self.n_init // world_size),
                      n_clusters=self.n_clusters,
                      max_iter=self.max_iter,
                      random_state=rank + 1,
                      verbose=0,
                      tol=self.tol,
                  ).fit(X.cpu().numpy())
                  self.inertia_ = _torch.Tensor([kmeans.inertia_]).to(
                      X.device
                  )
                  self.labels_ = (
                      _torch.from_numpy(kmeans.labels_)
                      .to(_torch.int)
                      .to(X.device)
                  )
                  self.cluster_centers_ = None
              else:
                  self.kmeans_pp(self.n_init, X, rank + 1)
  
              self.fit(X)
  
              bcast_rank = self.get_best_rank(
                  self.inertia_, _torch.argmin
              )
              if bcast_rank is not None:
                  _dist.broadcast(self.cluster_centers_, bcast_rank)
                  _dist.broadcast(self.labels_, bcast_rank)
  
              return self
  
          elif self.cluster_centers_ == "opt1d":
              nX, sort_order = _torch.sort(X, dim=0)
              nX = nX.cpu().numpy()
              rN = range(N)
  
              self.prefix_x = _np.cumsum(nX)
              self.prefix_x2 = _np.cumsum(_np.square(nX))
  
              new_cluster_cost = 0  # 2 * self.cost(0, N - 1, 0)
  
              num_D = self.n_clusters if self.verbose >= 2 else 2
  
              D = _np.full((num_D, N), _np.inf)
              D[0] = [self.cost(0, m, new_cluster_cost) for m in rN]
              T = _np.full((self.n_clusters, N), -1, dtype=int)
              T[0] = [0 for m in rN]
  
              opt_t_cost = D[0][-1]
              opt_n_clusters = 0
              for c in range(1, self.n_clusters):
                  if True:
  
                      def lookup(m, j):
                          return -(
                              D[(c - 1) % num_D][min(j - 1, m)]
                              + self.cost(j, m, new_cluster_cost)
                          )
  
                      R = self.smawk(rN, rN, lookup)
  
                      for k, v in R.items():
                          D[c % num_D][k] = -lookup(k, v)
                          T[c][k] = v
                  else:
                      for m in range(1, N):
                          for j in range(m):
                              cur_cost = D[(c - 1) % num_D][
                                  j
                              ] + self.cost(j + 1, m, new_cluster_cost)
                              if cur_cost < D[c % num_D][m]:
                                  D[c % num_D][m] = cur_cost
                                  T[c][m] = j + 1
  
                  if opt_t_cost > D[c % num_D][-1]:
                      opt_t_cost = D[c % num_D][-1]
                      opt_n_clusters = c
  
              P = []
              P = self.backtrace(
                  P, T, opt_n_clusters, T[opt_n_clusters][-1]
              )
              P.append(N)
  
              self.labels_ = []
              self.cluster_centers_ = []
              for i in range(len(P) - 1):
                  v = nX[P[i] : P[i + 1]]
                  if len(v):
                      self.labels_ += [
                          len(self.cluster_centers_)
                      ] * len(v)
                      self.cluster_centers_.append([_np.mean(v)])
  
              self.n_clusters = len(self.cluster_centers_)
              self.cluster_centers_ = _torch.from_numpy(
                  _np.array(self.cluster_centers_)
              ).to(device=X.device, dtype=X.dtype)
              min_error, self.labels_ = self.__class__.x_c_dist(
                  X, self.cluster_centers_
              ).min(dim=-1)
              self.inertia_ = min_error.sum()
  
      else:
          self.inertia_ = None
  
          for i in range(self.max_iter):
              self.cluster_centers_ = self.__class__.get_cluster_avg(
                  self.n_clusters, self.labels_, X
              )
  
              nan_centers = self.cluster_centers_.isnan()
              if nan_centers.any():
                  self.kmeans_pp(self.n_init, X, i)
                  continue
  
              self.x_c_dist = self.__class__.x_c_dist(
                  X, self.cluster_centers_
              )
              min_error, self.labels_ = self.x_c_dist.min(dim=-1)
              cur_inertia = min_error.sum()
  
              if (
                  self.error_bnd
                  and _torch.sqrt(cur_inertia / N) < self.error_bnd
              ):
                  unique, counts = _torch.unique(
                      self.labels_, return_counts=True
                  )
                  idx = unique[counts.argmin()]
  
                  reduce_cluster_centers_ = (
                      self.cluster_centers_.clone()
                  )
                  reduce_cluster_centers_[idx] = _np.nan
  
                  reduce_cluster_centers_ = reduce_cluster_centers_[
                      ~_torch.isnan(reduce_cluster_centers_)
                  ].view(-1, 1)
                  (
                      reduce_min_error,
                      reduce_labels_,
                  ) = self.__class__.x_c_dist(
                      X, reduce_cluster_centers_
                  ).min(
                      dim=-1
                  )
                  reduce_inertia = reduce_cluster_centers_.sum()
                  self.rmse_error = _torch.sqrt(reduce_inertia / N)
  
                  if self.rmse_error < self.error_bnd:
                      self.cluster_centers_ = reduce_cluster_centers_
                      self.labels_ = reduce_labels_
                      self.n_clusters = len(self.cluster_centers_)
                      continue
  
              if (
                  self.inertia_ is None
                  or abs(self.inertia_ - cur_inertia) > self.tol
              ):
                  self.inertia_ = cur_inertia
              else:
                  self.inertia_ = cur_inertia
                  break
  
      return self
  
file_path: coremltools/coremltools/optimize/torch/palettization/_efficient_kmeans.py
imports: import numpy
  import torch
  import torch.distributed
label_desc: Use image node_id 7 for calling the _EfficientKMeans obj's underlying member method code with example usage: obj.fit(X) and returns: self, self, unknown