import torch


def topk_pool_with_region(x, region, k=32):
    """
    Args:
        x: (B, F, M, P)
        region: (B, M, P)
        k: top k

    Returns:
        (B, F, M, topk)
    """
    featdim = x.shape[1]
    _, region_idx = torch.topk(region, k=k, dim=2)  # (B, M, k)
    index = region_idx.unsqueeze(1).repeat(1, featdim, 1, 1)  # (B, F, M, k)
    pooled = torch.gather(x, dim=3, index=index)
    return pooled