def blockwise_compress(
        original_data: np.ndarray,
        mode: str,
        nbits: Optional[int],
        block_sizes: List[int],
        lut_function: Optional[Callable] = None,
        cluster_dim: int = 1,
        channel_axis: Optional[int] = None,
        num_kmeans_workers: int = 1,
    ) -> Optional["optimize_utils.LutParams"]:
        """
        Compress original_data into n-bit representation by 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.

        channel_axis: Only useful for vector palettization (cluster_dim > 1). If not provided, we
            will try to infer it from `block_sizes`.

        Returns None if the weight cannot be compressed (for example, the dim size on an axis is not
        divisible by the corresponding block_size).
        """
        # TODO (rdar://127342739): Support more general blockwise palettization.
        # As general blockwise palettization hasn't been supported yet, we try to infer channel axis
        # and channel group size from block_sizes, and use grouped channelwise palettization instead.
        channel_group_size = 0
        for axis, block_size in enumerate(block_sizes):
            if block_size != 0 and block_size != original_data.shape[axis]:
                if channel_axis is not None and channel_axis != axis:
                    raise NotImplementedError(
                        "General block-wise palettization is not supported. Please use "
                        "'per_grouped_channel' or 'per_tensor' for the 'granularity' in config."
                    )
                channel_axis = axis
                channel_group_size = block_size
        if channel_axis is None:
            if cluster_dim > 1:
                raise ValueError(
                    "Cannot infer channel axis, which is required for vector palettization."
                )
            # Per-tensor compression, just need to pick a dummy axis.
            channel_axis = 0

        return palettize_weights.grouped_channelwise_compress(
            original_data,
            mode,
            nbits,
            channel_axis,
            channel_group_size,
            lut_function,
            cluster_dim,
            num_kmeans_workers,
        )

        Compress original_data into n-bit representation by 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.

        channel_axis: Only useful for vector palettization (cluster_dim > 1). If not provided, we
            will try to infer it from `block_sizes`.

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