import torch


def sample(data, uv):
    """Sample data (H, W, <C>) by uv (H, W, 2) (in pixels). """
    shape = data.shape
    # data from (H, W, <C>) to (1, C, H, W)
    data = data.reshape(data.shape[:2] + (-1,))
    data = torch.tensor(data).permute(2, 0, 1)[None, ...]
    # (H, W, 2) -> (1, H, W, 2)
    uv = torch.tensor(uv)[None, ...]

    H, W = shape[:2]
    # grid needs to be in [-1, 1] and (B, H, W, 2)
    size = torch.tensor((W, H), dtype=uv.dtype).view(1, 1, 1, -1)
    grid = (2 * uv / size - 1).to(data.dtype)
    tensor = torch.nn.functional.grid_sample(data, grid, padding_mode="border")
    # from (1, C, H, W) to (H, W, <C>)
    return tensor.permute(0, 2, 3, 1).reshape(shape).numpy()