def test_builder_eval_channelwise_lut(self, nbits):
        """
        Test channel-wise lut with first axis as channel axis (the first dim of lut has size > 1).

        indices = tensor<uint2, [6, 2]>>([2, 3, 3, 0, 1, 0, 3, 0, 2, 1, 0, 3])
        lut = tensor<int8, [2, 1, 4, 1]>([1, 5, 9, 13, 2, 10, 18, 26])

        It is effectively a 2-group 2-bit scalar palettization.
        The output shape would be [6, 2], which is the same as the indices shape.
        The output tensor values are:
        [[lut0[2]->9,  lut0[3]->13],
          [lut0[3]->13, lut0[0]->1],
          [lut0[1]->5,  lut0[0]->1],
          [lut1[3]->26, lut1[0]->2],
          [lut1[2]->18, lut1[1]->10],
          [lut1[0]->2,  lut1[3]->26]]
        """

        @mb.program(input_specs=[], opset_version=_IOS18_TARGET)
        def prog():
            if nbits == 1:
                indices = np.array([0, 1, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1]).reshape((6, 2))
                lut = np.array([1, 5, 9, 13]).reshape((2, 1, 2, 1)).astype(np.int8)
            else:
                indices = np.array([2, 3, 3, 0, 1, 0, 3, 0, 2, 1, 0, 3]).reshape((6, 2))
                lut = self._pad_lut_for_nbits_requirements(
                    np.array([1, 5, 9, 13, 2, 10, 18, 26]).reshape((2, 1, 4, 1)).astype(np.int8),
                    nbits=nbits,
                )
            indices_np_dtype = types.nptype_from_builtin(types.string_to_builtin(f"uint{nbits}"))
            indices = indices.astype(indices_np_dtype)
            return mb.constexpr_lut_to_dense(indices=indices, lut=lut)

        constexpr_lut_to_dense_op = prog.functions["main"].find_ops(
            op_type="constexpr_lut_to_dense"
        )[0]
        if nbits == 1:
            decompressed_res = np.array([1, 5, 5, 1, 5, 1, 13, 9, 13, 13, 9, 13])
        else:
            decompressed_res = np.array([9, 13, 13, 1, 5, 1, 26, 2, 18, 10, 2, 26])
        decompressed_res = decompressed_res.reshape((6, 2)).astype(np.int8)
        np.testing.assert_allclose(
            decompressed_res, constexpr_lut_to_dense_op.outputs[0].op.materialized_val_inference()
        )

        Test channel-wise lut with first axis as channel axis (the first dim of lut has size > 1).

        indices = tensor<uint2, [6, 2]>>([2, 3, 3, 0, 1, 0, 3, 0, 2, 1, 0, 3])
        lut = tensor<int8, [2, 1, 4, 1]>([1, 5, 9, 13, 2, 10, 18, 26])

        It is effectively a 2-group 2-bit scalar palettization.
        The output shape would be [6, 2], which is the same as the indices shape.
        The output tensor values are:
        [[lut0[2]->9,  lut0[3]->13],
          [lut0[3]->13, lut0[0]->1],
          [lut0[1]->5,  lut0[0]->1],
          [lut1[3]->26, lut1[0]->2],
          [lut1[2]->18, lut1[1]->10],
          [lut1[0]->2,  lut1[3]->26]]
        
coremltools/converters/mil/mil/ops/tests/iOS18/test_compression.py