class_docstr: Implements the beam search algorithm for decoding the most likely sequences.
  
  # Parameters
  
  end_index : `int`
      The index of the "stop" or "end" token in the target vocabulary.
  
  max_steps : `int`, optional (default = `50`)
      The maximum number of decoding steps to take, i.e. the maximum length
      of the predicted sequences.
  
  beam_size : `int`, optional (default = `10`)
      The width of the beam used.
  
  per_node_beam_size : `int`, optional (default = `beam_size`)
      The maximum number of candidates to consider per node, at each step in the search.
      If not given, this just defaults to `beam_size`. Setting this parameter
      to a number smaller than `beam_size` may give better results, as it can introduce
      more diversity into the search. See
      [*Beam Search Strategies for Neural Machine Translation*, Freitag and Al-Onaizan, 2017]
      (https://api.semanticscholar.org/CorpusID:2229477).
  
  sampler : `Sampler`, optional (default = `None`)
      An optional `Sampler` which is used to pick next candidate nodes and beams.
      If not specified, `DeterministicSampler` will be used, which just takes the
      `per_node_beam_size` most likely nodes and the `beam_size` most likely beams.
  
      Using the [`GumbelSampler`](#gumbelsampler), on the other hand, will give you
      [Stochastic Beam Search](https://api.semanticscholar.org/CorpusID:76662039).
  
  min_steps : `int`, optional (default = `None`)
      The minimum number of decoding steps to take, i.e. the minimum length of
      the predicted sequences. This does not include the start or end tokens. If `None`,
      no minimum is enforced.
  
  final_sequence_scorer : `FinalSequenceScorer`, optional (default = `None`)
      An optional `FinalSequenceScorer` which is used to score the final generated sequences.
      The output from this module is what is returned by the `search` method. If not
      specified, `SequenceLogProbabilityScorer` will be used, which scores the sequences
      by the sum of the token log probabilities.
  
  constraints: `List[Constraint]`, optional (default = `None`)
      An optional list of `Constraint`s which should be applied during beam search. If not
      provided, no constraints will be enforced.
  
  vocab: `Vocabulary`
      If `constraints` is not `None`, then `Vocabulary` will be passed to each constraint
      during its initialization. Having access to the vocabulary may be useful for certain
      contraints, e.g., to mask out invalid predictions during structured prediction.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "model", it gets specified as a top-level parameter, then is passed in to the model
      separately.
class_docstr_tok: ["Implements","the","beam","search","algorithm","for","decoding","the","most","likely","sequences",".","#","Parameters","end_index",":","`","int","`","The","index","of","the","``","stop","''","or","``","end","''","token","in","the","target","vocabulary",".","max_steps",":","`","int","`",",","optional","(","default","=","`","50","`",")","The","maximum","number","of","decoding","steps","to","take",",","i.e",".","the","maximum","length","of","the","predicted","sequences",".","beam_size",":","`","int","`",",","optional","(","default","=","`","10","`",")","The","width","of","the","beam","used",".","per_node_beam_size",":","`","int","`",",","optional","(","default","=","`","beam_size","`",")","The","maximum","number","of","candidates","to","consider","per","node",",","at","each","step","in","the","search",".","If","not","given",",","this","just","defaults","to","`","beam_size","`",".","Setting","this","parameter","to","a","number","smaller","than","`","beam_size","`","may","give","better","results",",","as","it","can","introduce","more","diversity","into","the","search",".","See","[","*","Beam","Search","Strategies","for","Neural","Machine","Translation","*",",","Freitag","and","Al-Onaizan",",","2017","]","(","https",":","\/\/api.semanticscholar.org\/CorpusID:2229477",")",".","sampler",":","`","Sampler","`",",","optional","(","default","=","`","None","`",")","An","optional","`","Sampler","`","which","is","used","to","pick","next","candidate","nodes","and","beams",".","If","not","specified",",","`","DeterministicSampler","`","will","be","used",",","which","just","takes","the","`","per_node_beam_size","`","most","likely","nodes","and","the","`","beam_size","`","most","likely","beams",".","Using","the","[","`","GumbelSampler","`","]","(","#","gumbelsampler",")",",","on","the","other","hand",",","will","give","you","[","Stochastic","Beam","Search","]","(","https",":","\/\/api.semanticscholar.org\/CorpusID:76662039",")",".","min_steps",":","`","int","`",",","optional","(","default","=","`","None","`",")","The","minimum","number","of","decoding","steps","to","take",",","i.e",".","the","minimum","length","of","the","predicted","sequences",".","This","does","not","include","the","start","or","end","tokens",".","If","`","None","`",",","no","minimum","is","enforced",".","final_sequence_scorer",":","`","FinalSequenceScorer","`",",","optional","(","default","=","`","None","`",")","An","optional","`","FinalSequenceScorer","`","which","is","used","to","score","the","final","generated","sequences",".","The","output","from","this","module","is","what","is","returned","by","the","`","search","`","method",".","If","not","specified",",","`","SequenceLogProbabilityScorer","`","will","be","used",",","which","scores","the","sequences","by","the","sum","of","the","token","log","probabilities",".","constraints",":","`","List","[","Constraint","]","`",",","optional","(","default","=","`","None","`",")","An","optional","list","of","`","Constraint","`","s","which","should","be","applied","during","beam","search",".","If","not","provided",",","no","constraints","will","be","enforced",".","vocab",":","`","Vocabulary","`","If","`","constraints","`","is","not","`","None","`",",","then","`","Vocabulary","`","will","be","passed","to","each","constraint","during","its","initialization",".","Having","access","to","the","vocabulary","may","be","useful","for","certain","contraints",",","e.g.",",","to","mask","out","invalid","predictions","during","structured","prediction",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","model","''",",","it","gets","specified","as","a","top-level","parameter",",","then","is","passed","in","to","the","model","separately","."]
code: def _search(
      self,
      start_predictions: torch.Tensor,
      start_state: StateType,
      step: StepFunctionTypeWithTimestep,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      batch_size = start_predictions.size()[0]
  
      # List of (batch_size, beam_size) tensors. One for each time step. Does not
      # include the start symbols, which are implicit.
      predictions: List[torch.Tensor] = []
  
      # List of (batch_size, beam_size) tensors. One for each time step. None for
      # the first.  Stores the index n for the parent prediction, i.e.
      # predictions[t-1][i][n], that it came from.
      backpointers: List[torch.Tensor] = []
  
      constraint_states = [
          constraint.init_state(batch_size)
          for constraint in self.constraints
      ]
  
      # Calculate the first timestep. This is done outside the main loop
      # because we are going from a single decoder input (the output from the
      # encoder) to the top `beam_size` decoder outputs. On the other hand,
      # within the main loop we are going from the `beam_size` elements of the
      # beam to `beam_size`^2 candidates from which we will select the top
      # `beam_size` elements for the next iteration.
      # shape: (batch_size, num_classes)
      start_class_log_probabilities, state = step(
          start_predictions, start_state, 0
      )
  
      num_classes = start_class_log_probabilities.size()[1]
  
      # Make sure `per_node_beam_size` is not larger than `num_classes`.
      if self.per_node_beam_size > num_classes:
          raise ConfigurationError(
              f"Target vocab size ({num_classes:d}) too small "
              f"relative to per_node_beam_size ({self.per_node_beam_size:d}).\n"
              f"Please decrease beam_size or per_node_beam_size."
          )
  
      sampler_state = self.sampler.init_state(
          start_class_log_probabilities, batch_size, num_classes
      )
  
      # Apply all constraints.
      if self.constraints:
          # shape: (batch_size, 1, num_classes)
          expanded_start_class_log_probabilities = (
              start_class_log_probabilities.unsqueeze(1)
          )
          for constraint, constraint_state in zip(
              self.constraints, constraint_states
          ):
              expanded_start_class_log_probabilities = constraint.apply(
                  constraint_state,
                  expanded_start_class_log_probabilities,
              )
          start_class_log_probabilities = (
              expanded_start_class_log_probabilities.squeeze(1)
          )
  
      # Prevent selecting the end symbol if there is any min_steps constraint
      if self.min_steps >= 1:
          start_class_log_probabilities[
              :, self._end_index
          ] = min_value_of_dtype(start_class_log_probabilities.dtype)
  
      # Get the initial predicted classed and their log probabilities.
      # shape: (batch_size, beam_size), (batch_size, beam_size)
      (
          start_top_log_probabilities,
          start_predicted_classes,
          sampler_state,
      ) = self.sampler.sample_beams(
          start_class_log_probabilities, self.beam_size, sampler_state
      )
  
      if (
          self.beam_size == 1
          and (start_predicted_classes == self._end_index).all()
      ):
          warnings.warn(
              "Empty sequences predicted. You may want to increase the beam size or ensure "
              "your step function is working properly.",
              RuntimeWarning,
          )
          return (
              start_predicted_classes.unsqueeze(-1),
              start_top_log_probabilities,
          )
  
      # The log probabilities for the last time step.
      # shape: (batch_size, beam_size)
      last_log_probabilities = start_top_log_probabilities
  
      # shape: [(batch_size, beam_size)]
      predictions.append(start_predicted_classes)
  
      # Log probability tensor that mandates that the end token is selected.
      # shape: (batch_size * beam_size, num_classes)
      log_probs_after_end = start_class_log_probabilities.new_full(
          (batch_size * self.beam_size, num_classes),
          min_value_of_dtype(start_class_log_probabilities.dtype),
      )
      log_probs_after_end[:, self._end_index] = 0.0
  
      # Set the same state for each element in the beam.
      self._update_initial_state(state, batch_size)
  
      for i, constraint in enumerate(self.constraints):
          constraint_states[i] = constraint.update_state(
              constraint_states[i], start_predicted_classes
          )
  
      for timestep in range(self.max_steps - 1):
          # shape: (batch_size * beam_size,)
          last_predictions = predictions[-1].reshape(
              batch_size * self.beam_size
          )
  
          # If every predicted token from the last step is `self._end_index`,
          # then we can stop early.
          if (last_predictions == self._end_index).all():
              break
          # Take a step. This get the predicted log probs of the next classes
          # and updates the state.
          # shape: (batch_size * beam_size, num_classes)
          class_log_probabilities, state = step(
              last_predictions, state, timestep + 1
          )
  
          # Apply all constraints.
          if self.constraints:
              # shape: (batch_size, beam_size, num_classes)
              reshaped_class_log_probabilities = (
                  class_log_probabilities.view(
                      batch_size, self.beam_size, -1
                  )
              )
              for constraint, constraint_state in zip(
                  self.constraints, constraint_states
              ):
                  reshaped_class_log_probabilities = constraint.apply(
                      constraint_state, reshaped_class_log_probabilities
                  )
              # shape: (batch_size * beam_size, num_classes)
              class_log_probabilities = (
                  reshaped_class_log_probabilities.view(
                      batch_size * self.beam_size, -1
                  )
              )
  
          # The `timestep`-th iteration of the for loop is generating the `timestep + 2`-th token
          # of the sequence (because `timestep` is 0-indexed and we generated the first token
          # before the for loop). Here we block the end index if the search is not allowed to
          # terminate on this iteration.
          if timestep + 2 <= self.min_steps:
              class_log_probabilities[
                  :, self._end_index
              ] = min_value_of_dtype(class_log_probabilities.dtype)
  
          # shape: (batch_size * beam_size, num_classes)
          last_predictions_expanded = last_predictions.unsqueeze(
              -1
          ).expand(batch_size * self.beam_size, num_classes)
  
          # Here we are finding any beams where we predicted the end token in
          # the previous timestep and replacing the distribution with a
          # one-hot distribution, forcing the beam to predict the end token
          # this timestep as well.
          # shape: (batch_size * beam_size, num_classes)
          cleaned_log_probabilities = torch.where(
              last_predictions_expanded == self._end_index,
              log_probs_after_end,
              class_log_probabilities,
          )
  
          # shape (both): (batch_size * beam_size, per_node_beam_size)
          (
              top_log_probabilities,
              predicted_classes,
              sampler_state,
          ) = self.sampler.sample_nodes(
              cleaned_log_probabilities,
              self.per_node_beam_size,
              sampler_state,
          )
  
          # Here we expand the last log probabilities to (batch_size * beam_size, per_node_beam_size)
          # so that we can add them to the current log probs for this timestep.
          # This lets us maintain the log probability of each element on the beam.
          # shape: (batch_size * beam_size, per_node_beam_size)
          expanded_last_log_probabilities = (
              last_log_probabilities.unsqueeze(2)
              .expand(
                  batch_size, self.beam_size, self.per_node_beam_size
              )
              .reshape(
                  batch_size * self.beam_size, self.per_node_beam_size
              )
          )
  
          # shape: (batch_size * beam_size, per_node_beam_size)
          summed_top_log_probabilities = (
              top_log_probabilities + expanded_last_log_probabilities
          )
  
          # shape: (batch_size, beam_size * per_node_beam_size)
          reshaped_summed = summed_top_log_probabilities.reshape(
              batch_size, self.beam_size * self.per_node_beam_size
          )
  
          # shape: (batch_size, beam_size * per_node_beam_size)
          reshaped_predicted_classes = predicted_classes.reshape(
              batch_size, self.beam_size * self.per_node_beam_size
          )
  
          # Keep only the top `beam_size` beam indices.
          # shape (both): (batch_size, beam_size)
          (
              restricted_beam_log_probs,
              restricted_beam_indices,
              sampler_state,
          ) = self.sampler.sample_beams(
              reshaped_summed, self.beam_size, sampler_state
          )
  
          # Use the beam indices to extract the corresponding classes.
          # shape: (batch_size, beam_size)
          restricted_predicted_classes = (
              reshaped_predicted_classes.gather(
                  1, restricted_beam_indices
              )
          )
  
          predictions.append(restricted_predicted_classes)
  
          # shape: (batch_size, beam_size)
          last_log_probabilities = restricted_beam_log_probs
  
          # The beam indices come from a `beam_size * per_node_beam_size` dimension where the
          # indices with a common ancestor are grouped together. Hence
          # dividing by per_node_beam_size gives the ancestor. (Note that this is integer
          # division as the tensor is a LongTensor.)
          # shape: (batch_size, beam_size)
          backpointer = torch.divide(
              restricted_beam_indices,
              self.per_node_beam_size,
              rounding_mode="trunc",
          )
          backpointers.append(backpointer)
  
          # Keep only the pieces of the state tensors corresponding to the
          # ancestors created this iteration.
          self._update_state(state, backpointer)
  
          for i, constraint in enumerate(self.constraints):
              constraint_states[i] = constraint.update_state(
                  constraint_states[i],
                  restricted_predicted_classes,
                  last_backpointer=backpointer,
              )
  
      # Warn about "-inf" log probabilities if not using any constraints (negligible
      # log probabilities are expected when using constraints).
      if not self.constraints and (
          not torch.isfinite(last_log_probabilities).all()
          or (
              last_log_probabilities
              == min_value_of_dtype(last_log_probabilities.dtype)
          ).any()
      ):
          warnings.warn(
              "Negligible log probabilities encountered ('-inf' or equivalent). "
              "Some final sequences may not make sense. "
              "This can happen when the beam size is larger than the number of valid (non-zero "
              "probability) transitions that the step function produces.",
              RuntimeWarning,
          )
  
      reconstructed_predictions = self._reconstruct_sequences(
          predictions, backpointers
      )
  
      # shape: (batch_size, beam_size, max_steps)
      all_predictions = torch.cat(
          list(reversed(reconstructed_predictions)), 2
      )
  
      # Calculate the final sequence scores
      # shape: (batch_size, beam_size)
      final_scores = self.final_sequence_scorer.score(
          all_predictions, last_log_probabilities, self._end_index
      )
  
      # Sort the sequences based on the final scores so the best scoring
      # sequence is at index 0
      sorted_final_scores, sorted_indices = torch.sort(
          final_scores, dim=1, descending=True
      )
      sorted_all_predictions = torch.gather(
          all_predictions,
          1,
          sorted_indices.unsqueeze(-1).expand_as(all_predictions),
      )
  
      return sorted_all_predictions, sorted_final_scores
  
imports: from inspect import signature
  from typing import Any, List, Callable, Tuple, Dict, cast, TypeVar, Optional
  import copy
  import warnings
  import torch
  from allennlp.common import Lazy, Registrable
  from allennlp.common.checks import ConfigurationError
  from allennlp.data import Vocabulary
  from allennlp.nn.util import min_value_of_dtype
label_desc: Use image node_id 4 for calling the BeamSearch obj's underlying member method code with example usage: obj._search(start_predictions, start_state, step) and returns: sorted_all_predictions,  sorted_final_scores, start_predicted_classes,  start_top_log_probabilities