import torch


def interpolate(feat, uv):
    """
    Args:
        feat (torch.Tensor): [B, C, H, W] image features
        uv (torch.Tensor): [B, 2, N] uv coordinates
            in the image plane, range [-1, 1]
    Returns:
        samples[:, :, :, 0] (torch.Tensor):
            [B, C, N] image features at the uv coordinates
    """
    if uv.shape[-1] != 2:
        uv = uv.transpose(1, 2)  # [B, N, 2]
    uv = uv.unsqueeze(2)  # [B, N, 1, 2]
    # NOTE: for newer PyTorch, it seems that training
    # results are degraded due to implementation diff in F.grid_sample
    # for old versions, simply remove the aligned_corners argument.
    if int(torch.__version__.split('.')[1]) < 4:
        samples = torch.nn.functional.grid_sample(feat, uv)  # [B, C, N, 1]
    else:
        samples = torch.nn.functional.grid_sample(
            feat, uv, align_corners=True)  # [B, C, N, 1]
    return samples[:, :, :, 0]