human
I wrote the below code

    batch_size = tf.shape(initial_ids)[0]
    initial_log_probs = tf.constant([([0.0] + ([(- float('inf'))] * (beam_size - 1)))])
    alive_log_probs = tf.tile(initial_log_probs, [batch_size, 1])
    alive_seq = tf.tile(tf.expand_dims(initial_ids, 1), [1, beam_size])
    alive_seq = tf.expand_dims(alive_seq, 2)
    finished_seq = tf.zeros(tf.shape(alive_seq), tf.int32)
    finished_scores = (tf.ones([batch_size, beam_size]) * (- INF))
    finished_flags = tf.zeros([batch_size, beam_size], tf.bool)
    def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq, curr_scores, curr_finished):
        'Given    sequences    and    scores,    will    gather    the    top    k=beam    size    sequences.\n\n                Args:\n                        finished_seq:    Current    finished    sequences.\n                                [batch_size,    beam_size,    current_decoded_length]\n                        finished_scores:    scores    for    each    of    these    sequences.\n                                [batch_size,    beam_size]\n                        finished_flags:    finished    bools    for    each    of    these    sequences.\n                                [batch_size,    beam_size]\n                        curr_seq:    current    topk    sequence    that    has    been    grown    by    one    position.\n                                [batch_size,    beam_size,    current_decoded_length]\n                        curr_scores:    scores    for    each    of    these    sequences.    [batch_size,    beam_size]\n                        curr_finished:    Finished    flags    for    each    of    these    sequences.\n                                [batch_size,    beam_size]\n                Returns:\n                        Tuple    of\n                                (Topk    sequences    based    on    scores,\n                                    log    probs    of    these    sequences,\n                                    Finished    flags    of    these    sequences)\n          ...

What's a good function header?
assistant
#Here's a good function header.

def beam_search(symbols_to_logits_fn, initial_ids, beam_size, decode_length, vocab_size, alpha, eos_id=EOS_ID):