import torch


def select_yx(featmap, y, x):
    """
    Select x, y coordinates from feature map.
    Args size:
        featmap: (B, C, H, W)
        x: (B, C)
        y: (B, C)
    """
    assert featmap.shape[:2] == x.shape == y.shape, 'X, Y coordinates should match.'
    x = torch.clamp(x, 0, featmap.shape[-1] - 1)
    y = torch.clamp(y, 0, featmap.shape[-2] - 1)
    b, c, h, w = featmap.shape
    y = y.view(b, c, 1, 1).repeat(1, 1, 1, w)
    featmap = torch.gather(featmap, -2, y.long())
    x = x.view(b, c, 1, 1)
    featmap = torch.gather(featmap, -1, x.long()) 
    return featmap.squeeze(-1).squeeze(-1)