def find_indices_for_lut(
    data: np.ndarray, lut: np.ndarray, vector_axis: Optional[int] = None
) -> np.ndarray:
    """
    Given a data and a look-up-table (LUT), find the closest indices in LUT that each element in
    data correspond to. It's the reverse process of "Given a LUT and indices, produce data using
    indices to fetch elements in LUT".

    Note the elements in data may not exactly match the elements in lut due to numerical instability.
    So we use fuzzy match to find the closest one instead of doing exact match.

    Parameters
    - data: Arbitrary numpy array.
    - lut: [block_num1, ..., 2**nbits, vector_size]. LUT's rank is K + 2, where K is the rank of data.
        Each dimension of data should be divisible by each corresponding dimension of the LUT.
        e.g., when data's shape is [2, 3, 4], the first three elements in lut's shape is [1, 1, 2],
        it means that there are two lookup tables over the last axis, and each of them have their
        own LUT values. See details in the iOS18 `constexpr_lut_to_dense` op.
    - vector_axis: Only effective when lut's last dim (vector_size) > 1. It denotes which axis the
        vector is along.
    """
    if len(lut.shape) != len(data.shape) + 2:
        raise ValueError(
            "The lut's rank should be data's rank + 2. See constexpr_lut_to_dense op definition."
        )
    if lut.shape[-1] > 1:
        if vector_axis is None:
            raise ValueError("The vector_axis must be provided for vector palettization.")
        if not len(data.shape) > vector_axis >= -len(data.shape):
            raise ValueError(f"Invalid vector_axis ({vector_axis})")
        if vector_axis < 0:
            vector_axis += len(data.shape)
        vector_size = lut.shape[-1]
        if data.shape[vector_axis] % vector_size != 0:
            raise ValueError(
                f"The data dim on {vector_axis}th axis ({data.shape[vector_axis]}) "
                f"must be divisible by vector_size ({vector_size})"
            )
        data = reshape_weight_for_vector_lut(data, vector_size, vector_axis)

    repeated_lut_shape = data.shape
    if lut.shape[-1] > 1:
        repeated_lut_shape = repeated_lut_shape[:-1]
    repeated_lut = repeat_data_as(lut, repeated_lut_shape)

    if lut.shape[-1] == 1:
        # For scalar palettization, we can simply find the closest value for each element.
        indices = np.argmin(
            np.abs(np.expand_dims(data, axis=-1) - np.squeeze(repeated_lut, axis=-1)), axis=-1
        )
    else:
        # For vector palettization, find the closest vector by Euclidean distance.
        dist = np.linalg.norm(np.expand_dims(data, axis=-2) - repeated_lut, axis=-1)
        indices = np.argmin(dist, axis=-1)

    nbits = int(math.log2(lut.shape[-2]))
    indices = indices.astype(types.nptype_from_builtin(types.string_to_builtin(f"uint{nbits}")))
    return indices

    Given a data and a look-up-table (LUT), find the closest indices in LUT that each element in
    data correspond to. It's the reverse process of "Given a LUT and indices, produce data using
    indices to fetch elements in LUT".

    Note the elements in data may not exactly match the elements in lut due to numerical instability.
    So we use fuzzy match to find the closest one instead of doing exact match.

    Parameters
    - data: Arbitrary numpy array.
    - lut: [block_num1, ..., 2**nbits, vector_size]. LUT's rank is K + 2, where K is the rank of data.
        Each dimension of data should be divisible by each corresponding dimension of the LUT.
        e.g., when data's shape is [2, 3, 4], the first three elements in lut's shape is [1, 1, 2],
        it means that there are two lookup tables over the last axis, and each of them have their
        own LUT values. See details in the iOS18 `constexpr_lut_to_dense` op.
    - vector_axis: Only effective when lut's last dim (vector_size) > 1. It denotes which axis the
        vector is along.
    
coremltools/optimize/_utils.py