def grouped_channelwise_compress(
        original_data: np.ndarray,
        mode: str,
        nbits: Optional[int],
        channel_axis: int,
        channel_group_size: int,
        lut_function: Optional[Callable] = None,
        cluster_dim: int = 1,
        num_kmeans_workers: int = 1,
    ) -> Optional["optimize_utils.LutParams"]:
        """
        Compress original_data into n-bit representation by grouped channelwise palettization.

        Supported nbits: 1, 2, 3, 4, 6, 8
        Supported mode: KMEANS, UNIFORM, UNIQUE, CUSTOM

        block_sizes: Each element is the block size on corresponding axis for original_data.

        cluster_dim: Dimension of each cluster centroid, which is the length of each element in the
            lookup table.

        Returns None if the weight cannot be compressed (for example, the dim size on an axis is not
        divisible by the corresponding channel_group_size).
        """
        if not isinstance(original_data, np.ndarray):
            raise ValueError(f"Only numpy arrays are supported, but got {type(original_data)}")
        if nbits is not None and nbits not in palettize_weights._SUPPORTED_NBITS:
            raise ValueError(
                f"Invalid nbits. Support {palettize_weights._SUPPORTED_NBITS}, but got {nbits}"
            )
        data_rank = len(original_data.shape)
        if not (-data_rank <= channel_axis < data_rank):
            raise ValueError(
                "Invalid channel_axis. Should be in range "
                f"[{-data_rank}, {data_rank}), but got {channel_axis}"
            )

        if channel_axis < 0:
            channel_axis += len(original_data.shape)

        channel_num = original_data.shape[channel_axis]
        if channel_group_size == 0:
            channel_group_size = channel_num
        if channel_num % channel_group_size != 0:
            logger.warning(
                f"Can't perform palettization: The number of channels at {channel_axis}th axis "
                f"({channel_num}) is not divisible by channel_group_size ({channel_group_size})."
            )
            return None
        channel_group_num = channel_num // channel_group_size

        if channel_group_size % cluster_dim != 0:
            logger.warning(
                f"Can't perform palettization: The channel_group_size at {channel_axis}th axis "
                f"({channel_group_size}) is not divisible by cluster_dim ({cluster_dim})."
            )
            return None

        if channel_axis != 0:
            original_data = np.swapaxes(original_data, 0, channel_axis)
        grouped_channel_data = np.split(original_data, channel_group_num, axis=0)

        # As the channel axis has been swapped to 0th axis, use 0 for vector_axis.
        vector_axis = 0

        # If mode is UNIQUE, infer nbits from the number of unique values in each group.
        if mode.upper() == "UNIQUE":
            try:
                for per_group_data in grouped_channel_data:
                    per_group_nbits = palettize_weights._get_nbits_for_unique_mode(
                        per_group_data, palettize_weights._SUPPORTED_NBITS, cluster_dim, vector_axis
                    )
                    # Pick the largest per-channel nbits to be used as the nbits for the whole op.
                    if nbits is None or per_group_nbits > nbits:
                        nbits = per_group_nbits
            except ValueError as e:
                logger.warning(f"Can't perform palettization:{e}")
                return None

        # The subprocesses have overhead, so only use it for expensive computations (k-means).
        if mode.upper() == "KMEANS" and num_kmeans_workers > 1:
            if palettize_weights._compress_pool is None:
                palettize_weights._compress_pool = Pool(processes=num_kmeans_workers)
                atexit.register(lambda: palettize_weights._compress_pool.terminate())
            lut, indices = zip(
                *palettize_weights._compress_pool.starmap(
                    palettize_weights._get_lut_and_indices,
                    zip(
                        grouped_channel_data,
                        repeat(mode),
                        repeat(nbits),
                        repeat(lut_function),
                        repeat(cluster_dim),
                        repeat(vector_axis),
                    ),
                )
            )
        else:
            lut, indices = zip(
                *[
                    palettize_weights._get_lut_and_indices(
                        per_channel_group_data, mode, nbits, lut_function, cluster_dim, vector_axis
                    )
                    for per_channel_group_data in grouped_channel_data
                ]
            )

        lut = np.stack(lut, axis=0)
        indices = np.stack(indices, axis=0)

        if mode.upper() == "CUSTOM":
            # The custom lut_function provided by users should have nbits info.
            # The current `lut` has shape [group_num, lut_entry_num, Optional[cluster_dim]].
            nbits = int(np.ceil(np.log2(lut.shape[1])))

        # The lut and indices from `_get_lut_and_indices` is flattened. The desired result should be
        # `lut` with shape [channel_group_num, palette_num], and `indices` with same shape as the
        # original_data.
        palette_num = 2**nbits
        indices_target_shape = list(original_data.shape)
        if cluster_dim > 1:
            indices_target_shape[vector_axis] //= cluster_dim
        indices = indices.reshape(indices_target_shape)
        lut_target_shape = [1] * (len(original_data.shape) + 2)
        lut_target_shape[0] = channel_group_num
        lut_target_shape[-1] = cluster_dim
        lut_target_shape[-2] = palette_num
        lut = lut.reshape(lut_target_shape)

        if channel_axis != 0:
            lut = np.swapaxes(lut, 0, channel_axis)
            indices = np.swapaxes(indices, 0, channel_axis)

        indices_np_dtype = types.nptype_from_builtin(types.string_to_builtin(f"uint{nbits}"))
        return optimize_utils.LutParams(
            indices.astype(indices_np_dtype), lut, None if cluster_dim == 1 else channel_axis
        )

        Compress original_data into n-bit representation by grouped channelwise palettization.

        Supported nbits: 1, 2, 3, 4, 6, 8
        Supported mode: KMEANS, UNIFORM, UNIQUE, CUSTOM

        block_sizes: Each element is the block size on corresponding axis for original_data.

        cluster_dim: Dimension of each cluster centroid, which is the length of each element in the
            lookup table.

        Returns None if the weight cannot be compressed (for example, the dim size on an axis is not
        divisible by the corresponding channel_group_size).
        
coremltools/optimize/coreml/_quantization_passes.py