def test_auto_pick_channel_axis_palettizer(self):
        """
        Check the right output channel axis is picked for granularity='per_grouped_channel'.
        """
        global_config = cto.coreml.OpPalettizerConfig(
            mode="kmeans",
            nbits=4,
            granularity="per_grouped_channel",
            group_size=1,
            weight_threshold=2000,
        )
        config = cto.coreml.OptimizationConfig()
        config.set_global(global_config)
        compressor = quantization.palettize_weights(config=config)

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

        # For conv, the output channel-axis is 0.
        conv_lut = prog.find_ops(op_type="constexpr_lut_to_dense")[0].lut
        assert conv_lut.shape[0] == 90
        assert conv_lut.shape[1] == 1
        # For linear, the output channel-axis is 0.
        linear_lut = prog.find_ops(op_type="constexpr_lut_to_dense")[1].lut
        assert linear_lut.shape[0] == 70
        assert linear_lut.shape[1] == 1
        # For matmul with transpose_y=False, the output channel-axis is -1.
        matmul_lut = prog.find_ops(op_type="constexpr_lut_to_dense")[2].lut
        assert matmul_lut.shape == (1, 1, 1, 35, 16, 1)
        # For conv_transpose, the output channel-axis is -2.
        conv_transpose_lut = prog.find_ops(op_type="constexpr_lut_to_dense")[3].lut
        assert conv_transpose_lut.shape[0] == 1
        assert conv_transpose_lut.shape[1] == 4

        Check the right output channel axis is picked for granularity='per_grouped_channel'.
        
coremltools/test/optimize/coreml/test_passes.py