(initial_sequence, initial_state, generate_step_fn, num_steps,
                beam_size, branch_factor, steps_per_iteration)
return beam_entry.sequence, beam_entry.state, beam_entry.score
Generates a sequence using beam search.

  Initially, the beam is filled with `beam_size` copies of the initial sequence.

  Each iteration, the beam is pruned to contain only the `beam_size` event
  sequences with highest score. Then `branch_factor` new event sequences are
  generated for each sequence in the beam. These new sequences are formed by
  extending each sequence in the beam by `steps_per_iteration` steps. So between
  a branching and a pruning phase, there will be `beam_size` * `branch_factor`
  active event sequences.

  After the final iteration, the single sequence in the beam with highest
  likelihood will be returned.

  The `generate_step_fn` function operates on lists of sequences + states +
  scores rather than single sequences. This is to allow for the possibility of
  batching.

  Args:
    initial_sequence: The initial sequence, a Python list-like object.
    initial_state: The state corresponding to the initial sequence, with any
        auxiliary information needed for extending the sequence.
    generate_step_fn: A function that takes three parameters: a list of
        sequences, a list of states, and a list of scores, all of the same size.
        The function should generate a single step for each of the sequences and
        return the extended sequences, updated states, and updated (total)
        scores, as three lists.
    num_steps: The integer length in steps of the final sequence, after
        generation.
    beam_size: The integer beam size to use.
    branch_factor: The integer branch factor to use.
    steps_per_iteration: The integer number of steps to take per iteration.

  Returns:
    A tuple containing a) the highest-scoring sequence as computed by the beam
    search, b) the state corresponding to this sequence, and c) the score of
    this sequence.
Generates a sequence using beam search.
Generates a sequence using beam search .
def beam_search(initial_sequence, initial_state, generate_step_fn, num_steps,
                beam_size, branch_factor, steps_per_iteration):
  """Generates a sequence using beam search.

  Initially, the beam is filled with `beam_size` copies of the initial sequence.

  Each iteration, the beam is pruned to contain only the `beam_size` event
  sequences with highest score. Then `branch_factor` new event sequences are
  generated for each sequence in the beam. These new sequences are formed by
  extending each sequence in the beam by `steps_per_iteration` steps. So between
  a branching and a pruning phase, there will be `beam_size` * `branch_factor`
  active event sequences.

  After the final iteration, the single sequence in the beam with highest
  likelihood will be returned.

  The `generate_step_fn` function operates on lists of sequences + states +
  scores rather than single sequences. This is to allow for the possibility of
  batching.

  Args:
    initial_sequence: The initial sequence, a Python list-like object.
    initial_state: The state corresponding to the initial sequence, with any
        auxiliary information needed for extending the sequence.
    generate_step_fn: A function that takes three parameters: a list of
        sequences, a list of states, and a list of scores, all of the same size.
        The function should generate a single step for each of the sequences and
        return the extended sequences, updated states, and updated (total)
        scores, as three lists.
    num_steps: The integer length in steps of the final sequence, after
        generation.
    beam_size: The integer beam size to use.
    branch_factor: The integer branch factor to use.
    steps_per_iteration: The integer number of steps to take per iteration.

  Returns:
    A tuple containing a) the highest-scoring sequence as computed by the beam
    search, b) the state corresponding to this sequence, and c) the score of
    this sequence.
  """
  sequences = [copy.deepcopy(initial_sequence) for _ in range(beam_size)]
  states = [copy.deepcopy(initial_state) for _ in range(beam_size)]
  scores = [0] * beam_size

  beam_entries = [BeamEntry(sequence, state, score)
                  for sequence, state, score
                  in zip(sequences, states, scores)]

  # Choose the number of steps for the first iteration such that subsequent
  # iterations can all take the same number of steps.
  first_iteration_num_steps = (num_steps - 1) % steps_per_iteration + 1

  beam_entries = _generate_branches(
      beam_entries, generate_step_fn, branch_factor, first_iteration_num_steps)

  num_iterations = (num_steps -
                    first_iteration_num_steps) // steps_per_iteration

  for _ in range(num_iterations):
    beam_entries = _prune_branches(beam_entries, k=beam_size)
    beam_entries = _generate_branches(
        beam_entries, generate_step_fn, branch_factor, steps_per_iteration)

  # Prune to the single best beam entry.
  beam_entry = _prune_branches(beam_entries, k=1)[0]

  return beam_entry.sequence, beam_entry.state, beam_entry.score
def beam_search ( initial_sequence , initial_state , generate_step_fn , num_steps , beam_size , branch_factor , steps_per_iteration ) : sequences = [ copy . deepcopy ( initial_sequence ) for _ in range ( beam_size ) ] states = [ copy . deepcopy ( initial_state ) for _ in range ( beam_size ) ] scores = [ 0 ] * beam_size beam_entries = [ BeamEntry ( sequence , state , score ) for sequence , state , score in zip ( sequences , states , scores ) ] # Choose the number of steps for the first iteration such that subsequent # iterations can all take the same number of steps. first_iteration_num_steps = ( num_steps - 1 ) % steps_per_iteration + 1 beam_entries = _generate_branches ( beam_entries , generate_step_fn , branch_factor , first_iteration_num_steps ) num_iterations = ( num_steps - first_iteration_num_steps ) // steps_per_iteration for _ in range ( num_iterations ) : beam_entries = _prune_branches ( beam_entries , k = beam_size ) beam_entries = _generate_branches ( beam_entries , generate_step_fn , branch_factor , steps_per_iteration ) # Prune to the single best beam entry. beam_entry = _prune_branches ( beam_entries , k = 1 ) [ 0 ] return beam_entry . sequence , beam_entry . state , beam_entry . score