import torch


def interpolate_costs(cost_maps: torch.Tensor, p2D: torch.Tensor):
    """Interpolate cost values in dense cost maps.
    Args:
        * cost_maps: The dense [N x H x W] cost maps.
        * p2D: The [N x 2] keypoint coordinates.
    Returns:
        * costs: The [N] tensor of bilinearly interpolated costs.
    """
    assert len(cost_maps) == len(p2D), "Dimensions do not match"
    height, width = cost_maps.shape[-2:]
    scale = torch.tensor([width - 1, height - 1], dtype=p2D.dtype, device=p2D.device)
    p2D = ((p2D / scale) * 2 - 1).clamp(min=-2, max=2)
    costs = torch.nn.functional.grid_sample(
        cost_maps[:, None],
        p2D[:, None, None],
        mode="bilinear",
        align_corners=True,
        padding_mode="border",
    )
    return costs.view(len(p2D))