def init_partitions(self, parameters) -> None:
        """
        Method to initialize the partitions and set the k-means. Called during first iteration of palettization in the
        forward method of ``FakePalettize``.
        """
        if isinstance(self.centroids[0], _torch.Tensor):
            if self.centroids[0].device != parameters.device:
                self._move_centroids_and_labels_to_device(parameters.device, parameters.dtype)
            return
        with _torch.no_grad():
            num_partitions = len(self.partitions)
            numel_per_partition = max(self.partition_numel)
            if "nf" in self.centroids[0]:
                if self.n_clusters in NF_CLUSTER_SIZES and self.cluster_dim == 1:
                    nf_fit = "fit" in self.centroids[0]
                    for i, partition in enumerate(self.partitions):
                        bit = int(_math.log2(self.n_clusters))
                        sparse = bool(_math.log2(self.n_clusters) - bit)

                        self.centroids[i] = (
                            _generate_natural_float(bit=bit, sparse=sparse)
                            .to(parameters.device)
                            .to(parameters.dtype)
                            .view(-1, 1)
                        )

                        if nf_fit:
                            best_err = _torch.finfo(_torch.float).max
                            best_lambd = 1
                            best_retry = 0
                            best_thold = 10
                            up_down_hill = 0
                            lambd_list = [[1 + x / 100, 1 - x / 100] for x in range(99)]
                            lambd_list = [1] + [v for sublist in lambd_list for v in sublist]

                            cur_X = parameters[self.partitions[i]].view(-1, 1)

                            for cur_lambd in lambd_list:
                                if up_down_hill > best_thold and cur_lambd < 1:
                                    continue

                                if up_down_hill < -best_thold and cur_lambd > 1:
                                    continue

                                cur_lut = _torch.stack(
                                    [x.sign() * x.abs() ** (cur_lambd) for x in self.centroids[i]]
                                )
                                x_c_dist = _torch.cdist(cur_X, cur_lut.to(cur_X.dtype)).square()
                                cur_err = x_c_dist.min(-1).values.float().sum()

                                if best_err > cur_err:
                                    best_retry = 0
                                    best_err = cur_err
                                    best_lambd = cur_lambd
                                    if best_lambd > 1:
                                        up_down_hill += 1
                                    else:
                                        up_down_hill -= 1

                                elif best_retry > best_thold:
                                    break
                                else:
                                    best_retry += 1

                            self.centroids[i] = _torch.stack(
                                [x.sign() * x.abs() ** (best_lambd) for x in self.centroids[i]]
                            )
                    return

                self.centroids = ["auto"] * num_partitions

            for i in range(num_partitions):
                if self.centroids[i] == "auto":
                    # if auto then pick either init method
                    self.centroids[i] = (
                        "opt1d"
                        if (
                            numel_per_partition <= self.n_clusters
                            or numel_per_partition <= self.kmeans_opt1d_threshold
                        )
                        and self.cluster_dim == 1
                        else "kmeans++"
                    )

            if _dist.is_available() and _dist.is_initialized():
                distributed_world_size = _dist.get_world_size()
            else:
                distributed_world_size = 1
            if max(num_partitions, distributed_world_size) < self.kmeans_batch_threshold:
                for i, partition in enumerate(self.partitions):
                    kmeans = self.get_partition_kmeans(
                        parameters,
                        partition,
                        self.n_clusters,
                        self.labels[i],
                        self.enforce_zero[i],
                        max_iter=100,
                        init=self.centroids[i],
                        n_init=max(1, self.kmeans_n_init // distributed_world_size),
                    )
                    bcast_rank = _get_best_rank(kmeans.inertia_, _torch.argmin)
                    if bcast_rank:
                        _dist.broadcast(kmeans.cluster_centers_, bcast_rank)

                    self.centroids[i] = kmeans.cluster_centers_
                    self.labels[i] = None
            else:
                shard_list = _get_shard_list(num_partitions)
                centroids_list = [None] * distributed_world_size

                for i in range(distributed_world_size):
                    begin, end = shard_list[i], shard_list[i + 1]
                    current_rank = (
                        _dist.get_rank() if _dist.is_available() and _dist.is_initialized() else 0
                    )
                    if i == current_rank and begin < end:
                        for p in range(begin, end):
                            kmeans = self.get_partition_kmeans(
                                parameters,
                                self.partitions[p],
                                self.n_clusters,
                                self.labels[p],
                                self.enforce_zero[p],
                                max_iter=100,
                                init=self.centroids[p],
                                n_init=self.kmeans_n_init,
                            )
                            self.centroids[p] = kmeans.cluster_centers_

                        centroids_list[i] = _torch.stack(self.centroids[begin:end])
                    else:
                        centroids_list[i] = _torch.full(
                            [end - begin, self.n_clusters, self.cluster_dim],
                            float("nan"),
                            dtype=parameters.dtype,
                            device=parameters.device,
                        )

                if _dist.is_available() and _dist.is_initialized():
                    _dist.all_gather(centroids_list, centroids_list[_dist.get_rank()])
                centroids_list = [v for sublist in centroids_list for v in sublist]

                assert len(centroids_list) == num_partitions
                for p in range(num_partitions):
                    self.labels[p] = None
                    self.centroids[p] = centroids_list[p]

        Method to initialize the partitions and set the k-means. Called during first iteration of palettization in the
        forward method of ``FakePalettize``.
        
coremltools/optimize/torch/palettization/_partitioner.py