def search(self,
               initial_state: State,
               transition_function: TransitionFunction) -> Dict[int, List[State]]:
        """
        Parameters
        ----------
        initial_state : ``State``
            The starting state of our search.  This is assumed to be `batched`, and our beam search
            is batch-aware - we'll keep ``beam_size`` states around for each instance in the batch.
        transition_function : ``TransitionFunction``
            The ``TransitionFunction`` object that defines and scores transitions from one state to the
            next.

        Returns
        -------
        best_states : ``Dict[int, List[State]]``
            This is a mapping from batch index to the top states for that instance.
        """
        finished_states: Dict[int, List[State]] = defaultdict(list)
        states = [initial_state]
        step_num = 0
        while states:
            step_num += 1
            next_states: Dict[int, List[State]] = defaultdict(list)
            grouped_state = states[0].combine_states(states)
            allowed_actions = []
            for batch_index, action_history in zip(grouped_state.batch_indices,
                                                   grouped_state.action_history):
                allowed_actions.append(self._allowed_transitions[batch_index][tuple(action_history)])
            for next_state in transition_function.take_step(grouped_state,
                                                            max_actions=self._per_node_beam_size,
                                                            allowed_actions=allowed_actions):
                # NOTE: we're doing state.batch_indices[0] here (and similar things below),
                # hard-coding a group size of 1.  But, our use of `next_state.is_finished()`
                # already checks for that, as it crashes if the group size is not 1.
                batch_index = next_state.batch_indices[0]
                if next_state.is_finished():
                    finished_states[batch_index].append(next_state)
                else:
                    next_states[batch_index].append(next_state)
            states = []
            for batch_index, batch_states in next_states.items():
                # The states from the generator are already sorted, so we can just take the first
                # ones here, without an additional sort.
                if self._beam_size:
                    batch_states = batch_states[:self._beam_size]
                states.extend(batch_states)
        best_states: Dict[int, List[State]] = {}
        for batch_index, batch_states in finished_states.items():
            # The time this sort takes is pretty negligible, no particular need to optimize this
            # yet.  Maybe with a larger beam size...
            finished_to_sort = [(-state.score[0].item(), state) for state in batch_states]
            finished_to_sort.sort(key=lambda x: x[0])
            best_states[batch_index] = [state[1] for state in finished_to_sort[:self._beam_size]]
        return best_states
Parameters
        ----------
        initial_state : ``State``
            The starting state of our search.  This is assumed to be `batched`, and our beam search
            is batch-aware - we'll keep ``beam_size`` states around for each instance in the batch.
        transition_function : ``TransitionFunction``
            The ``TransitionFunction`` object that defines and scores transitions from one state to the
            next.

        Returns
        -------
        best_states : ``Dict[int, List[State]]``
            This is a mapping from batch index to the top states for that instance.