(beam_entries, generate_step_fn, branch_factor,
                       num_steps)
return [BeamEntry(sequence, state, score)
          for sequence, state, score
          in zip(all_sequences, all_states, all_scores)]
Performs a single iteration of branch generation for beam search.

  This method generates `branch_factor` branches for each sequence in the beam,
  where each branch extends the event sequence by `num_steps` steps (via calls
  to `generate_step_fn`). The resulting beam is returned.

  Args:
    beam_entries: A list of BeamEntry tuples, the current beam.
    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. The function may modify the sequences, states,
        and scores in place, but should also return the modified values.
    branch_factor: The integer branch factor to use.
    num_steps: The integer number of steps to take per branch.

  Returns:
    The updated beam, with `branch_factor` times as many BeamEntry tuples.
Performs a single iteration of branch generation for beam search.
Performs a single iteration of branch generation for beam search .
def _generate_branches(beam_entries, generate_step_fn, branch_factor,
                       num_steps):
  """Performs a single iteration of branch generation for beam search.

  This method generates `branch_factor` branches for each sequence in the beam,
  where each branch extends the event sequence by `num_steps` steps (via calls
  to `generate_step_fn`). The resulting beam is returned.

  Args:
    beam_entries: A list of BeamEntry tuples, the current beam.
    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. The function may modify the sequences, states,
        and scores in place, but should also return the modified values.
    branch_factor: The integer branch factor to use.
    num_steps: The integer number of steps to take per branch.

  Returns:
    The updated beam, with `branch_factor` times as many BeamEntry tuples.
  """
  if branch_factor > 1:
    branched_entries = beam_entries * branch_factor
    all_sequences = [copy.deepcopy(entry.sequence)
                     for entry in branched_entries]
    all_states = [copy.deepcopy(entry.state) for entry in branched_entries]
    all_scores = [entry.score for entry in branched_entries]
  else:
    # No need to make copies if there's no branching.
    all_sequences = [entry.sequence for entry in beam_entries]
    all_states = [entry.state for entry in beam_entries]
    all_scores = [entry.score for entry in beam_entries]

  for _ in range(num_steps):
    all_sequences, all_states, all_scores = generate_step_fn(
        all_sequences, all_states, all_scores)

  return [BeamEntry(sequence, state, score)
          for sequence, state, score
          in zip(all_sequences, all_states, all_scores)]
def _generate_branches ( beam_entries , generate_step_fn , branch_factor , num_steps ) : if branch_factor > 1 : branched_entries = beam_entries * branch_factor all_sequences = [ copy . deepcopy ( entry . sequence ) for entry in branched_entries ] all_states = [ copy . deepcopy ( entry . state ) for entry in branched_entries ] all_scores = [ entry . score for entry in branched_entries ] else : # No need to make copies if there's no branching. all_sequences = [ entry . sequence for entry in beam_entries ] all_states = [ entry . state for entry in beam_entries ] all_scores = [ entry . score for entry in beam_entries ] for _ in range ( num_steps ) : all_sequences , all_states , all_scores = generate_step_fn ( all_sequences , all_states , all_scores ) return [ BeamEntry ( sequence , state , score ) for sequence , state , score in zip ( all_sequences , all_states , all_scores ) ]