code: def _gather_topk_beams(nested, score_or_log_prob, batch_size, beam_size):
    
    _, topk_indexes = tf.nn.top_k(score_or_log_prob, k=beam_size)
    return _gather_beams(nested, topk_indexes, batch_size, beam_size)
docs: Gather top beams from nested structure.