def index_function_wrapper(fun):
    """
    Modify a multi-index function to accept multi-dimensional arrays of
    multi-indices.
    """

    def new_fun(inds):
        shape = inds.shape
        new_inds = inds.reshape(-1, shape[-1])
        out = fun(new_inds)
        return out.reshape(shape[:-1])

    return new_fun