import torch.nn.functional as F

def bilinear_interpolate(fmap, pos):
    
    assert fmap.size(0) == pos.size(0), "Batch size of fmap and pos must be the same"
    grid = pos.unsqueeze(-2).unsqueeze(-2)  # Shape will be (N, 1, 1, 2)
    grid = grid.to(fmap.device)
    interpolated_value = F.grid_sample(fmap, grid, mode="bilinear", padding_mode="zeros", align_corners=True)

    return interpolated_value.squeeze(-1).squeeze(-1)
   
def bilinear_interpolate_samples(fmap, pos):

    assert fmap.size(0) == pos.size(0), 'Batch size of fmap and pos must be the same'
    
    grid = pos.unsqueeze(-2)
    grid = grid.to(fmap.device)
    interpolated_value = F.grid_sample(fmap, grid, mode="bilinear", padding_mode="zeros", align_corners=True)
    scores = interpolated_value.squeeze(-1).permute(0,2,1)
    
    return scores