import torch


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

    Returns:
        (B, F, M, topk)
    """
    _, idx = torch.topk(x, k=k, dim=3)  # (B, F, M, k)
    pooled = torch.gather(x, dim=3, index=idx)
    return pooled