def _get_lut_and_indices(
        val: np.ndarray,
        mode: str,
        nbits: Optional[int],
        lut_function: Optional[Callable],
        cluster_dim: int = 1,
        vector_axis: Optional[int] = None,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Calculate look-up-table (LUT) and indices."""

        def compress_kmeans(val, nbits, cluster_dim, vector_axis):
            lut, indices = _get_kmeans_lookup_table_and_weight(
                nbits, val, force_kmeans1d=False, cluster_dim=cluster_dim, vector_axis=vector_axis
            )
            lut = lut.astype(val.dtype)
            indices = indices.astype(np.uint8)
            return lut, indices

        def compress_uniform(val, nbits):
            val = val.flatten()
            val_min = np.amin(val)
            val_max = np.amax(val)
            scale = (val_max - val_min) / ((1 << nbits) - 1)
            indices = np.round(((val - val_min) / (val_max - val_min)) * ((1 << nbits) - 1)).astype(
                np.uint8
            )
            lut = np.array(range(0, 1 << nbits)) * scale + val_min
            lut = lut.astype(val.dtype)
            return lut, indices

        def compress_unique(val, nbits, cluster_dim, vector_axis):
            if nbits is None:
                nbits = palettize_weights._get_nbits_for_unique_mode(
                    val,
                    palettize_weights._SUPPORTED_NBITS,
                    cluster_dim,
                    vector_axis,
                )

            if cluster_dim > 1:
                val = optimize_utils.reshape_weight_for_vector_lut(val, cluster_dim, vector_axis)

            val = val.reshape((-1, cluster_dim))
            unique_vals, unique_inverse = np.unique(val, axis=0, return_inverse=True)
            lut = np.zeros((1 << nbits, cluster_dim))
            lut[: len(unique_vals)] = unique_vals
            indices = unique_inverse
            indices = indices.flatten()

            if cluster_dim == 1:
                # Squeeze the last dim to make behaviors back compatible with scalar palettization.
                lut = lut.squeeze(-1)

            return lut.astype(val.dtype), indices.astype(np.uint8)

        if mode == "KMEANS":
            lut, indices = compress_kmeans(val, nbits, cluster_dim, vector_axis)
        elif mode == "UNIFORM":
            if cluster_dim > 1:
                raise NotImplementedError(
                    "Vector palettization (cluster_dim > 1) doesn't support UNIFORM mode."
                )
            lut, indices = compress_uniform(val, nbits)
        elif mode == "UNIQUE":
            lut, indices = compress_unique(val, nbits, cluster_dim, vector_axis)
        else:
            if mode != "CUSTOM":
                raise AssertionError(f"Invalid mode {mode}")
            lut, indices = lut_function(val)

        return lut, indices
Calculate look-up-table (LUT) and indices.
coremltools/optimize/coreml/_quantization_passes.py