import torch


def rerank_beams_v2(beams, scores):
    """
        beams: [batch, beam, len]
        scores: [batch, beam], negative logprobs
    """
    _, reranked_ids = torch.sort(scores, dim=1, descending=False)  # [batch, beam]
    tiled_reranked_ids = reranked_ids.unsqueeze(-1).repeat(1, 1, beams.size(-1))
    reranked_beams = torch.gather(beams, dim=1,
                                  index=tiled_reranked_ids)
    return reranked_beams