code: 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
  
imports: import collections
  import copy
  import heapq
label_desc: Use image node_id 3 for calling a global function with example usage: beam_search(initial_sequence, initial_state, generate_step_fn, num_steps, beam_size, branch_factor, steps_per_iteration) and returns: beam_entry,  beam_entry,  beam_entry