def test_builder_eval_vector_lut(self, compute_unit, backend, vector_axis):
        """
        Test vector lut on different axis.

        indices = [
                    [
                       [4, 8], -> group 0
                       [10, 13], -> group 0
                       [24, 5], -> group 1
                       [6, 9] -> group 1
                    ],
                    [
                       [13, 31], -> group 0
                       [17, 7], -> group 0
                       [2, 8], -> group 1
                       [3, 1] -> group 1
                    ]
                  ]
        """
        def build():
            return mb.constexpr_lut_to_dense(
                indices=np.array([4, 8, 10, 13, 24, 5, 6, 9, 13, 31, 17, 7, 2, 8, 3, 1])
                .reshape((2, 4, 2))
                .astype(np.uint8),
                lut=self._generate_lut(shape=(1, 2, 1, 256, 3)),
                vector_axis=vector_axis,
            )

        prog = mb.program(input_specs=[], opset_version=_IOS18_TARGET)(build)
        constexpr_lut_to_dense_op = prog.functions["main"].find_ops(
            op_type="constexpr_lut_to_dense"
        )[0]
        if vector_axis == 0:
            decompressed_res = (
                np.array(
                    [
                        12,
                        24,
                        30,
                        39,
                        127,
                        127,
                        127,
                        127,
                        13,
                        25,
                        31,
                        40,
                        127,
                        127,
                        127,
                        127,
                        14,
                        26,
                        32,
                        41,
                        127,
                        127,
                        127,
                        127,
                        39,
                        93,
                        51,
                        21,
                        127,
                        127,
                        127,
                        127,
                        40,
                        94,
                        52,
                        22,
                        127,
                        127,
                        127,
                        127,
                        41,
                        95,
                        53,
                        23,
                        127,
                        127,
                        127,
                        127,
                    ]
                )
                .reshape((2 * 3, 4, 2))
                .astype(np.int8)
            )
        elif vector_axis == 1:
            decompressed_res = (
                np.array(
                    [
                        12,
                        24,
                        13,
                        25,
                        14,
                        26,
                        30,
                        39,
                        31,
                        40,
                        32,
                        41,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        39,
                        93,
                        40,
                        94,
                        41,
                        95,
                        51,
                        21,
                        52,
                        22,
                        53,
                        23,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                    ]
                )
                .reshape((2, 4 * 3, 2))
                .astype(np.int8)
            )
        else:
            decompressed_res = (
                np.array(
                    [
                        12,
                        13,
                        14,
                        24,
                        25,
                        26,
                        30,
                        31,
                        32,
                        39,
                        40,
                        41,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        39,
                        40,
                        41,
                        93,
                        94,
                        95,
                        51,
                        52,
                        53,
                        21,
                        22,
                        23,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                        127,
                    ]
                )
                .reshape((2, 4, 2 * 3))
                .astype(np.int8)
            )
        np.testing.assert_allclose(
            decompressed_res, constexpr_lut_to_dense_op.outputs[0].op.materialized_val_inference()
        )

        run_compare_builder(
            build,
            {},
            input_values={},
            expected_output_types=decompressed_res.shape + (types.fp32,),
            expected_outputs=decompressed_res,
            compute_unit=compute_unit,
            backend=backend,
        )

        Test vector lut on different axis.

        indices = [
                    [
                       [4, 8], -> group 0
                       [10, 13], -> group 0
                       [24, 5], -> group 1
                       [6, 9] -> group 1
                    ],
                    [
                       [13, 31], -> group 0
                       [17, 7], -> group 0
                       [2, 8], -> group 1
                       [3, 1] -> group 1
                    ]
                  ]
        
coremltools/converters/mil/mil/ops/tests/iOS18/test_compression.py