import torch


def interpolate_irreg_grid(interpfunc, pos):
    """Interpolate the  funtion

    Args:
        interpfunc (callable): function to interpolate the data points
        pos (torch.tensor): positions of the walkers Nbatch x 3*Nelec

    Returns:
        torch.tensor: interpolated values of the function evaluated at pos
    """

    nbatch, nelec, ndim = pos.shape[0], pos.shape[1]//3, 3
    return torch.as_tensor(interpfunc(pos.reshape(nbatch, nelec, ndim)))