def test_auto_pick_channel_axis_quantizer(self):
        """
        Check the right output channel axis is picked for block-wise quantization.
        """
        global_config = cto.coreml.OpLinearQuantizerConfig(
            mode="LINEAR",
            dtype="int4",
            granularity="per_block",
            block_size=2,
            weight_threshold=2000,
        )
        linear_config = cto.coreml.OpLinearQuantizerConfig(
            mode="LINEAR_SYMMETRIC",
            dtype="int4",
            granularity="per_block",
            block_size=9,
            weight_threshold=100,
        )
        config = cto.coreml.OptimizationConfig()
        config.set_global(global_config)
        config.set_op_name("linear", linear_config)
        compressor = quantization.linear_quantize_weights(config=config)

        prog = self._get_test_program_3()
        compressor.apply(prog)

        blockwise_ops = prog.find_ops(op_type="constexpr_blockwise_shift_scale")
        # For conv, input channel axis is 1, output channel axis is 0.
        # The original weight shape is [90, 30, 2, 2], the scale's second dim is 30 / 2 = 15.
        assert blockwise_ops[0].scale.shape == (90, 15, 1, 1)
        # For linear, input channel axis is 1, output channel axis is 0.
        # The original weight shape is [70, 81], the scale's second dim is 81 / 9 = 9.
        assert blockwise_ops[1].scale.shape == (70, 9)
        # For matmul (transpose_y=False), input channel axis is -2, output channel axis is -1.
        # The original weight shape is [2, 1, 70, 35], the scale's third dim is 70 / 2 = 35.
        assert blockwise_ops[2].scale.shape == (1, 1, 35, 35)
        # For conv_transpose, input channel axis is 0, output channel axis is 1.
        # The original weight shape is [30, 4, 21, 10], the scale's first dim is 30 / 2 = 15.
        assert blockwise_ops[3].scale.shape == (15, 4, 1, 1)

        Check the right output channel axis is picked for block-wise quantization.
        
coremltools/test/optimize/coreml/test_passes.py