def lut_to_dense(
    indices: np.ndarray,
    lut: np.ndarray,
    vector_axis: Optional[int] = None,
) -> np.ndarray:
    """Decompress the lut with indices to dense (decompressed) output."""
    palette_num = lut.shape[-2]
    vector_size = lut.shape[-1]

    output_shape = list(indices.shape)
    if vector_size > 1:
        if vector_axis < 0:
            vector_axis += len(output_shape)
        # Swap the vector axis to the last dim to vectorize computations.
        indices = np.swapaxes(indices, -1, vector_axis)
        lut = np.swapaxes(lut, -3, vector_axis)
        output_shape = list(indices.shape)
        output_shape[-1] *= vector_size

    flattened_indices = indices.astype(np.uint8).flatten()
    repeated_lut = repeat_data_as(lut, indices.shape).reshape(
        len(flattened_indices), palette_num, vector_size
    )

    output = repeated_lut[np.arange(repeated_lut.shape[0]), flattened_indices.astype(np.int32)]
    output = output.reshape(output_shape)
    if vector_size > 1:
        output = np.swapaxes(output, -1, vector_axis)

    return output
Decompress the lut with indices to dense (decompressed) output.
coremltools/optimize/_utils.py