def _update_clustered_block_weight(
        cls,
        block_weight: _torch.Tensor,
        block_mask: _torch.Tensor,
        depalett_block_weight: _torch.Tensor,
        new_weight: _torch.Tensor,
        config: KMeansConfig,
        block_idx: int,
    ):
        """
        Write back clustered weight in new weight.
        """
        block_weight_flatten = block_weight.flatten()

        if block_mask is not None:
            new_block_weight = block_weight_flatten.clone()
            new_block_weight[block_mask] = depalett_block_weight
            new_block_weight = new_block_weight.reshape(block_weight.shape)
        else:
            if config.axis == 1 or config.cluster_dim == 1:
                new_block_weight = depalett_block_weight.reshape(block_weight.shape)
            else:
                # need to reshape back for cluster_dim > 1 and axis = 0
                new_block_weight = depalett_block_weight.reshape(
                    block_weight.shape[1], block_weight.shape[0]
                ).transpose(0, 1)

        if config.axis == 0:
            new_weight[block_idx : block_idx + config.block_size, :] = new_block_weight
        else:
            new_weight[:, block_idx : block_idx + config.block_size] = new_block_weight

        Write back clustered weight in new weight.
        
coremltools/optimize/torch/_utils/k_means.py