def test_grouped_channelwise_equivalent_to_blockwise(
        self, nbits, channel_axis, channel_group_size
    ):
        """The grouped channelwise palettization could be expressed as general blockwise."""
        original_data = np.random.randint(low=-256, high=256, size=(16, 16, 2, 2)).astype(
            np.float32
        )

        params_grouped_channelwise = quantization.palettize_weights.grouped_channelwise_compress(
            original_data, "UNIFORM", nbits, channel_axis, channel_group_size
        )
        decompressed_grouped_channelwise = quantization.palettize_weights.decompress(
            params_grouped_channelwise
        )

        block_sizes = [0] * len(original_data.shape)
        block_sizes[channel_axis] = channel_group_size
        params_blockwise = quantization.palettize_weights.blockwise_compress(
            original_data, "UNIFORM", nbits, block_sizes=block_sizes
        )
        decompressed_blockwise = quantization.palettize_weights.decompress(params_blockwise)

        np.testing.assert_allclose(
            np.sort(params_grouped_channelwise.lut, axis=None),
            np.sort(params_blockwise.lut, axis=None),
        )
        np.testing.assert_allclose(decompressed_grouped_channelwise, decompressed_blockwise)
The grouped channelwise palettization could be expressed as general blockwise.
coremltools/test/optimize/coreml/test_passes.py