    def search(
        self,
        start_predictions: torch.Tensor,
        start_state: StateType,
        step: StepFunctionType,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Given a starting state and a step function, apply beam search to find the
        most likely target sequences.

        !!! Note
            If your step function returns `-inf` for some log probabilities
            (like if you're using a masked log-softmax) then some of the "best"
            sequences returned may also have `-inf` log probability. Specifically
            this happens when the beam size is smaller than the number of actions
            with finite log probability (non-zero probability) returned by the step function.
            Therefore if you're using a mask you may want to check the results from `search`
            and potentially discard sequences with non-finite log probability.

        # Parameters

        start_predictions : `torch.Tensor`
            A tensor containing the initial predictions with shape `(batch_size,)`.
            Usually the initial predictions are just the index of the "start" token
            in the target vocabulary.

        start_state : `StateType`
            The initial state passed to the `step` function. Each value of the state dict
            should be a tensor of shape `(batch_size, *)`, where `*` means any other
            number of dimensions.

        step : `StepFunctionType`
            A function that is responsible for computing the next most likely tokens,
            given the current state and the predictions from the last time step.
            The function should accept two or three arguments:

            - a tensor of shape `(group_size,)` representing the index of the predicted
            tokens from the last time step,
            - the current state, a `StateType`, and
            - optionally, the timestep, an `int`.

            The `group_size` will be `batch_size * beam_size`, except in the initial
            step, for which it will just be `batch_size`.

            The function is expected to return a tuple, where the first element
            is a tensor of shape `(group_size, target_vocab_size)` containing
            the log probabilities of the tokens for the next step, and the second
            element is the updated state. The tensor in the state should have shape
            `(group_size, *)`, where `*` means any other number of dimensions.

        # Returns

        `Tuple[torch.Tensor, torch.Tensor]`
            Tuple of `(predictions, final_scores)`, where `predictions`
            has shape `(batch_size, beam_size, max_steps)` and `final_scores`
            has shape `(batch_size, beam_size)`.
        """
        step_signature = signature(step)
        if len(step_signature.parameters) < 3:
            # If the step function we're given does not take the time step argument, wrap it
            # in one that does.
            old_step = cast(StepFunctionTypeNoTimestep, step)

            def new_step(
                last_predictions: torch.Tensor, state: Dict[str, torch.Tensor], time_step: int
            ):
                return old_step(last_predictions, state)

            return self._search(start_predictions, start_state, new_step)
        else:
            return self._search(
                start_predictions, start_state, cast(StepFunctionTypeWithTimestep, step)
            )
Given a starting state and a step function, apply beam search to find the
most likely target sequences.

!!! Note
If your step function returns `-inf` for some log probabilities
(like if you're using a masked log-softmax) then some of the "best"
sequences returned may also have `-inf` log probability. Specifically
this happens when the beam size is smaller than the number of actions
with finite log probability (non-zero probability) returned by the step function.
Therefore if you're using a mask you may want to check the results from `search`
and potentially discard sequences with non-finite log probability.

# Parameters

start_predictions : `torch.Tensor`
A tensor containing the initial predictions with shape `(batch_size,)`.
Usually the initial predictions are just the index of the "start" token
in the target vocabulary.

start_state : `StateType`
The initial state passed to the `step` function. Each value of the state dict
should be a tensor of shape `(batch_size, *)`, where `*` means any other
number of dimensions.

step : `StepFunctionType`
A function that is responsible for computing the next most likely tokens,
given the current state and the predictions from the last time step.
The function should accept two or three arguments:

- a tensor of shape `(group_size,)` representing the index of the predicted
tokens from the last time step,
- the current state, a `StateType`, and
- optionally, the timestep, an `int`.

The `group_size` will be `batch_size * beam_size`, except in the initial
step, for which it will just be `batch_size`.

The function is expected to return a tuple, where the first element
is a tensor of shape `(group_size, target_vocab_size)` containing
the log probabilities of the tokens for the next step, and the second
element is the updated state. The tensor in the state should have shape
`(group_size, *)`, where `*` means any other number of dimensions.

# Returns

`Tuple[torch.Tensor, torch.Tensor]`
Tuple of `(predictions, final_scores)`, where `predictions`
has shape `(batch_size, beam_size, max_steps)` and `final_scores`
has shape `(batch_size, beam_size)`.
allennlp/nn/beam_search.py