    def sample_beams(
        self,
        log_probs: torch.Tensor,
        beam_size: int,
        state: StateType,
    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
        """
        Returns the beams with the highest perturbed log probabilities.
        """
        # shape (log_probs): (batch_size, beam_size * per_node_beam_size)

        batch_size = log_probs.size()[0]

        # shape: (batch_size * beam_size, per_node_beam_size)
        G_phi_S = state["G_phi_S"]

        # shape: (batch_size, beam_size * per_node_beam_size)
        G_phi_S = G_phi_S.reshape_as(log_probs)

        # shape (both): (batch_size, beam_size)
        G_phi_S_new, selected_indices = torch.topk(G_phi_S, beam_size, dim=-1)

        # shape: (batch_size, beam_size)
        selected_log_probs = log_probs.gather(1, selected_indices)

        # Now sort the selected beams by their true log prob.
        # shape (all): (batch_size, beam_size)
        selected_log_probs, sort_indices = selected_log_probs.sort(dim=-1, descending=True)
        selected_indices = selected_indices.gather(1, sort_indices)
        G_phi_S_new = G_phi_S_new.gather(1, sort_indices)

        # shape: (batch_size * beam_size,)
        G_phi_S_new = G_phi_S_new.reshape(batch_size * beam_size)

        # shape: (batch_size * beam_size,)
        phi_S = selected_log_probs.reshape(batch_size * beam_size)

        return selected_log_probs, selected_indices, {"G_phi_S": G_phi_S_new, "phi_S": phi_S}
Returns the beams with the highest perturbed log probabilities.
allennlp/nn/beam_search.py