def get_partition_kmeans(
        self,
        X,
        partition,
        n_clusters,
        labels,
        enforce_zero,
        max_iter,
        init,
        n_init=10,
    ) -> _EfficientKMeans:
        """
        Method to get kmeans for a particular partition.
        """
        cY, pad = _vectorize(X[partition], self.cluster_dim, self.vector_ch_axis)

        kmeans = _EfficientKMeans(
            n_clusters=n_clusters,
            init=init,
            labels=labels,
            n_init=n_init,
            max_iter=max_iter,
            error_bnd=self.kmeans_error_bnd,
        ).fit(cY)

        if enforce_zero:
            # fix zero
            zero_point = _torch.zeros_like(kmeans.cluster_centers_[0]).unsqueeze(0)
            zero_idx = _torch.argmin(
                _torch.cdist(kmeans.cluster_centers_.float(), zero_point.float())
            )

            # always put zero in the first
            temp = kmeans.cluster_centers_[0]
            kmeans.cluster_centers_[zero_idx] = temp
            kmeans.cluster_centers_[0] = zero_point

        return kmeans

        Method to get kmeans for a particular partition.
        
coremltools/optimize/torch/palettization/_partitioner.py