def infer_block_sizes(
    op: "Operation",
    op_config: Union[OpLinearQuantizerConfig, OpPalettizerConfig],
    weight_to_compress: np.ndarray,
    return_channel_axis: bool = False,
) -> Union[Optional[List[int]], Tuple[Optional[List[int]], Optional[int]]]:
    """
    Infer block size on each axis based on the op and compression config.

    For per-channel, the channel axis is auto-picked.
    For per-block, the input/output axis is auto-picked if block_size is int.
    See the docstring of OpLinearQuantizerConfig for more details.
    """
    input_channel_axis, output_channel_axis = select_input_output_channel_axis(op)

    if op_config.granularity == CompressionGranularity.PER_BLOCK and not isinstance(
        op_config.block_size, int
    ):
        if len(op_config.block_size) != len(weight_to_compress.shape):
            raise ValueError(
                "The block_size in config must has one element for each axis. However, for op "
                f"{op.name}, there are {len(op_config.block_size)} elements in block_size, "
                f"but there are {len(weight_to_compress.shape)} axes in the weight."
            )
        channel_axis_candidates = []
        for axis, (b_size, dim_size) in enumerate(
            zip(op_config.block_size, weight_to_compress.shape)
        ):
            if b_size != 0 and b_size != dim_size:
                channel_axis_candidates.append(axis)
        if len(channel_axis_candidates) == 1:
            # Set channel axis if we can infer it from block sizes; else just use the default one
            # inferred by op type.
            output_channel_axis = channel_axis_candidates[0]

        if return_channel_axis:
            return list(op_config.block_size), output_channel_axis
        return list(op_config.block_size)

    if input_channel_axis is None or output_channel_axis is None:
        if return_channel_axis:
            return None, output_channel_axis
        return None

    if (
        op_config.granularity == CompressionGranularity.PER_GROUPED_CHANNEL
        and op_config.channel_axis is not None
    ):
        output_channel_axis = op_config.channel_axis

    block_sizes = [0] * len(weight_to_compress.shape)
    if op_config.granularity == CompressionGranularity.PER_TENSOR:
        input_channel_block_size = 0
        output_channel_block_size = 0
    elif op_config.granularity == CompressionGranularity.PER_CHANNEL:
        input_channel_block_size = 0
        output_channel_block_size = 1
    elif op_config.granularity == CompressionGranularity.PER_GROUPED_CHANNEL:
        input_channel_block_size = 0
        output_channel_block_size = op_config.group_size
    else:
        assert op_config.granularity == CompressionGranularity.PER_BLOCK and isinstance(
            op_config.block_size, int
        )
        input_channel_block_size = op_config.block_size
        output_channel_block_size = 1

    if input_channel_axis < len(block_sizes):
        block_sizes[input_channel_axis] = input_channel_block_size
    if output_channel_axis < len(block_sizes):
        block_sizes[output_channel_axis] = output_channel_block_size

    if return_channel_axis:
        return block_sizes, output_channel_axis
    return block_sizes

    Infer block size on each axis based on the op and compression config.

    For per-channel, the channel axis is auto-picked.
    For per-block, the input/output axis is auto-picked if block_size is int.
    See the docstring of OpLinearQuantizerConfig for more details.
    
coremltools/optimize/_utils.py