class_docstr: A class containing all of the functions supporting generation, to be used as a mixin in [`TFPreTrainedModel`].
  
  The class exposes [`~generation.TFGenerationMixin.generate`], which can be used for:
      - *greedy decoding* by calling [`~generation.TFGenerationMixin.greedy_search`] if `num_beams=1` and
        `do_sample=False`
      - *contrastive search* by calling [`~generation.TFGenerationMixin.contrastive_search`] if `penalty_alpha>0` and
        `top_k>1`
      - *multinomial sampling* by calling [`~generation.TFGenerationMixin.sample`] if `num_beams=1` and
        `do_sample=True`
      - *beam-search decoding* by calling [`~generation.TFGenerationMixin.beam_search`] if `num_beams>1`
  
  You do not need to call any of the above methods directly. Pass custom parameter values to 'generate' instead. To
  learn more about decoding strategies refer to the [text generation strategies guide](../generation_strategies).
class_docstr_tok: ["A","class","containing","all","of","the","functions","supporting","generation",",","to","be","used","as","a","mixin","in","[","`","TFPreTrainedModel","`","]",".","The","class","exposes","[","`","~generation.TFGenerationMixin.generate","`","]",",","which","can","be","used","for",":","-","*","greedy","decoding","*","by","calling","[","`","~generation.TFGenerationMixin.greedy_search","`","]","if","`","num_beams=1","`","and","`","do_sample=False","`","-","*","contrastive","search","*","by","calling","[","`","~generation.TFGenerationMixin.contrastive_search","`","]","if","`","penalty_alpha",">","0","`","and","`","top_k",">","1","`","-","*","multinomial","sampling","*","by","calling","[","`","~generation.TFGenerationMixin.sample","`","]","if","`","num_beams=1","`","and","`","do_sample=True","`","-","*","beam-search","decoding","*","by","calling","[","`","~generation.TFGenerationMixin.beam_search","`","]","if","`","num_beams",">","1","`","You","do","not","need","to","call","any","of","the","above","methods","directly",".","Pass","custom","parameter","values","to","'generate","'","instead",".","To","learn","more","about","decoding","strategies","refer","to","the","[","text","generation","strategies","guide","]","(","..","\/generation_strategies",")","."]
docstr: Generates sequences for models with a language modeling head using beam search. If `do_sample` is `False`, uses
  a greedy approach, otherwise does multinomial sampling without replacement.
  
  Parameters:
      input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`):
          The sequence used as a prompt for the generation.
      do_sample (`bool`, *optional*, defaults to `False`):
          Whether or not to use sampling ; use greedy decoding otherwise.
      max_length (`int`, *optional*, defaults to 20):
          The maximum length of the sequence to be generated.
      pad_token_id (`int`, *optional*):
          The id of the *padding* token.
      eos_token_id (`Union[int, List[int]]`, *optional*):
          The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens.
      length_penalty (`float`, *optional*, defaults to 1.0):
          Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent
          to the sequence length, which in turn is used to divide the score of the sequence. Since the score is
          the log likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences,
          while `length_penalty` < 0.0 encourages shorter sequences.
      early_stopping (`bool` or `str`, *optional*, defaults to `False`):
          Controls the stopping condition for beam-based methods, like beam-search. It accepts the following
          values: `True`, where the generation stops as soon as there are `num_beams` complete candidates;
          `False`, where an heuristic is applied and the generation stops when is it very unlikely to find better
          candidates; `"never"`, where the beam search procedure only stops when there cannot be better
          candidates (canonical beam search algorithm).
      logits_processor (`[TFLogitsProcessorList]`, *optional*):
          An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsProcessor`]
          used to modify the prediction scores of the language modeling head applied at each generation step.
      logits_warper (`TFLogitsProcessorList`, *optional*):
          An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsWarper`]
          used to warp the prediction score distribution of the language modeling head applied before multinomial
          sampling at each generation step.
      num_return_sequences(`int`, *optional*, defaults to 1):
          The number of independently computed returned sequences for each element in the batch.
      output_attentions (`bool`, *optional*, defaults to `False`):
          Whether or not to return the attentions tensors of all attention layers. See `attentions` under
          returned tensors for more details.
      output_hidden_states (`bool`, *optional*, defaults to `False`):
          Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
          for more details.
      return_dict_in_generate (`bool`, *optional*, defaults to `False`):
          Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
      model_kwargs:
          Additional model specific kwargs will be forwarded to the `call` function of the model. If model is an
          encoder-decoder model the kwargs should include `encoder_outputs`.
  
  Return:
      [`~generation.TFBeamSearchDecoderOnlyOutput`], [`~generation.TFBeamSearchEncoderDecoderOutput`] or
      `tf.Tensor`: A `tf.Tensor` containing the generated tokens (default behaviour) or a
      [`~generation.TFBeamSearchDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
      `return_dict_in_generate=True` or a [`~generation.TFBeamSearchEncoderDecoderOutput`] if
      `model.config.is_encoder_decoder=True`.
  
  Examples:
  
  ```python
  >>> from transformers import (
  ...     AutoTokenizer,
  ...     TFAutoModelForSeq2SeqLM,
  ...     TFLogitsProcessorList,
  ...     TFMinLengthLogitsProcessor,
  ... )
  >>> import tensorflow as tf
  
  >>> tokenizer = AutoTokenizer.from_pretrained("t5-base")
  >>> model = TFAutoModelForSeq2SeqLM.from_pretrained("t5-base")
  
  >>> encoder_input_str = "translate English to German: How old are you?"
  >>> encoder_input_ids = tokenizer(encoder_input_str, return_tensors="tf").input_ids
  
  >>> # lets run beam search using 3 beams
  >>> num_beams = 3
  >>> # define decoder start token ids
  >>> input_ids = tf.ones((1, num_beams, 1), dtype=tf.int32)
  >>> input_ids = input_ids * model.generation_config.decoder_start_token_id
  
  >>> # add encoder_outputs to model keyword arguments
  >>> encoder_outputs = model.get_encoder()(encoder_input_ids, return_dict=True)
  >>> encoder_outputs.last_hidden_state = tf.repeat(
  ...     tf.expand_dims(encoder_outputs.last_hidden_state, axis=0), num_beams, axis=1
  ... )
  >>> model_kwargs = {"encoder_outputs": encoder_outputs}
  
  >>> # instantiate logits processors
  >>> logits_processor = TFLogitsProcessorList(
  ...     [TFMinLengthLogitsProcessor(5, eos_token_id=model.generation_config.eos_token_id)]
  ... )
  
  >>> outputs = model.beam_search(input_ids, logits_processor=logits_processor, **model_kwargs)
  >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
  ['Wie alt bist du?']
  ```
code: def beam_search(
      self,
      input_ids: tf.Tensor,
      do_sample: bool = False,
      max_length: Optional[int] = None,
      pad_token_id: Optional[int] = None,
      eos_token_id: Optional[int] = None,
      length_penalty: Optional[float] = None,
      early_stopping: Optional[Union[bool, str]] = None,
      logits_processor: Optional[TFLogitsProcessorList] = None,
      logits_warper: Optional[TFLogitsProcessorList] = None,
      num_return_sequences: Optional[int] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      output_scores: Optional[bool] = None,
      return_dict_in_generate: Optional[bool] = None,
      **model_kwargs,
  ) -> Union[TFBeamSearchOutput, TFBeamSampleOutput, tf.Tensor]:
      r"""
      Generates sequences for models with a language modeling head using beam search. If `do_sample` is `False`, uses
      a greedy approach, otherwise does multinomial sampling without replacement.
  
      Parameters:
          input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`):
              The sequence used as a prompt for the generation.
          do_sample (`bool`, *optional*, defaults to `False`):
              Whether or not to use sampling ; use greedy decoding otherwise.
          max_length (`int`, *optional*, defaults to 20):
              The maximum length of the sequence to be generated.
          pad_token_id (`int`, *optional*):
              The id of the *padding* token.
          eos_token_id (`Union[int, List[int]]`, *optional*):
              The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens.
          length_penalty (`float`, *optional*, defaults to 1.0):
              Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent
              to the sequence length, which in turn is used to divide the score of the sequence. Since the score is
              the log likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences,
              while `length_penalty` < 0.0 encourages shorter sequences.
          early_stopping (`bool` or `str`, *optional*, defaults to `False`):
              Controls the stopping condition for beam-based methods, like beam-search. It accepts the following
              values: `True`, where the generation stops as soon as there are `num_beams` complete candidates;
              `False`, where an heuristic is applied and the generation stops when is it very unlikely to find better
              candidates; `"never"`, where the beam search procedure only stops when there cannot be better
              candidates (canonical beam search algorithm).
          logits_processor (`[TFLogitsProcessorList]`, *optional*):
              An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsProcessor`]
              used to modify the prediction scores of the language modeling head applied at each generation step.
          logits_warper (`TFLogitsProcessorList`, *optional*):
              An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsWarper`]
              used to warp the prediction score distribution of the language modeling head applied before multinomial
              sampling at each generation step.
          num_return_sequences(`int`, *optional*, defaults to 1):
              The number of independently computed returned sequences for each element in the batch.
          output_attentions (`bool`, *optional*, defaults to `False`):
              Whether or not to return the attentions tensors of all attention layers. See `attentions` under
              returned tensors for more details.
          output_hidden_states (`bool`, *optional*, defaults to `False`):
              Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
              for more details.
          return_dict_in_generate (`bool`, *optional*, defaults to `False`):
              Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
          model_kwargs:
              Additional model specific kwargs will be forwarded to the `call` function of the model. If model is an
              encoder-decoder model the kwargs should include `encoder_outputs`.
  
      Return:
          [`~generation.TFBeamSearchDecoderOnlyOutput`], [`~generation.TFBeamSearchEncoderDecoderOutput`] or
          `tf.Tensor`: A `tf.Tensor` containing the generated tokens (default behaviour) or a
          [`~generation.TFBeamSearchDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
          `return_dict_in_generate=True` or a [`~generation.TFBeamSearchEncoderDecoderOutput`] if
          `model.config.is_encoder_decoder=True`.
  
      Examples:
  
      ```python
      >>> from transformers import (
      ...     AutoTokenizer,
      ...     TFAutoModelForSeq2SeqLM,
      ...     TFLogitsProcessorList,
      ...     TFMinLengthLogitsProcessor,
      ... )
      >>> import tensorflow as tf
  
      >>> tokenizer = AutoTokenizer.from_pretrained("t5-base")
      >>> model = TFAutoModelForSeq2SeqLM.from_pretrained("t5-base")
  
      >>> encoder_input_str = "translate English to German: How old are you?"
      >>> encoder_input_ids = tokenizer(encoder_input_str, return_tensors="tf").input_ids
  
      >>> # lets run beam search using 3 beams
      >>> num_beams = 3
      >>> # define decoder start token ids
      >>> input_ids = tf.ones((1, num_beams, 1), dtype=tf.int32)
      >>> input_ids = input_ids * model.generation_config.decoder_start_token_id
  
      >>> # add encoder_outputs to model keyword arguments
      >>> encoder_outputs = model.get_encoder()(encoder_input_ids, return_dict=True)
      >>> encoder_outputs.last_hidden_state = tf.repeat(
      ...     tf.expand_dims(encoder_outputs.last_hidden_state, axis=0), num_beams, axis=1
      ... )
      >>> model_kwargs = {"encoder_outputs": encoder_outputs}
  
      >>> # instantiate logits processors
      >>> logits_processor = TFLogitsProcessorList(
      ...     [TFMinLengthLogitsProcessor(5, eos_token_id=model.generation_config.eos_token_id)]
      ... )
  
      >>> outputs = model.beam_search(input_ids, logits_processor=logits_processor, **model_kwargs)
      >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
      ['Wie alt bist du?']
      ```"""
  
      def flatten_beam_dim(tensor, batch_axis=0):
          """Flattens the first two dimensions of a non-scalar array."""
          shape = shape_list(tensor)
          return tf.reshape(
              tensor,
              shape[:batch_axis]
              + [shape[batch_axis] * shape[batch_axis + 1]]
              + shape[batch_axis + 2 :],
          )
  
      def unflatten_beam_dim(tensor, num_beams, batch_axis=0):
          """Unflattens the first, flat batch*beam dimension of a non-scalar array."""
          shape = shape_list(tensor)
          return tf.reshape(
              tensor,
              shape[:batch_axis]
              + [-1, num_beams]
              + shape[batch_axis + 1 :],
          )
  
      # 1. init beam_search values
      logits_processor = (
          logits_processor
          if logits_processor is not None
          else TFLogitsProcessorList()
      )
      logits_warper = (
          logits_warper
          if logits_warper is not None
          else TFLogitsProcessorList()
      )
  
      max_length = (
          max_length
          if max_length is not None
          else self.generation_config.max_length
      )
      pad_token_id = (
          pad_token_id
          if pad_token_id is not None
          else self.generation_config.pad_token_id
      )
      eos_token_id = (
          eos_token_id
          if eos_token_id is not None
          else self.generation_config.eos_token_id
      )
      if isinstance(eos_token_id, int):
          eos_token_id = [eos_token_id]
      num_return_sequences = (
          num_return_sequences
          if num_return_sequences is not None
          else self.generation_config.num_return_sequences
      )
  
      output_attentions = (
          output_attentions
          if output_attentions is not None
          else self.generation_config.output_attentions
      )
      output_hidden_states = (
          output_hidden_states
          if output_hidden_states is not None
          else self.generation_config.output_hidden_states
      )
      output_scores = (
          output_scores
          if output_scores is not None
          else self.generation_config.output_scores
      )
      return_dict_in_generate = (
          return_dict_in_generate
          if return_dict_in_generate is not None
          else self.generation_config.return_dict_in_generate
      )
  
      length_penalty = (
          length_penalty
          if length_penalty is not None
          else self.generation_config.length_penalty
      )
      early_stopping = (
          early_stopping
          if early_stopping is not None
          else self.generation_config.early_stopping
      )
  
      use_cache = model_kwargs.pop(
          "use_cache", self.generation_config.use_cache
      )
      use_xla = not tf.executing_eagerly()
      # TODO (Joao): fix cache format or find programatic way to detect cache index
      # GPT2 and other models has a slightly different cache structure, with a different batch axis
      model_name = (
          str(self.decoder)
          if "EncoderDecoder" in str(self)
          else str(self)
      )
      cache_batch_axis = (
          1
          if any(
              model_prefix in model_name
              for model_prefix in ("TFGPT2", "TFCTRL")
          )
          else 0
      )
      # some models, like XLNet, need more than the last token in the presence of past_key_values
      needs_full_input = "use_mems" in set(
          inspect.signature(
              self.prepare_inputs_for_generation
          ).parameters.keys()
      )
  
      # 2. init `attentions`, `hidden_states`, and `scores` tuples
      all_scores = (
          [] if (return_dict_in_generate and output_scores) else None
      )
      decoder_attentions = (
          []
          if (return_dict_in_generate and output_attentions)
          else None
      )
      cross_attentions = (
          []
          if (return_dict_in_generate and output_attentions)
          else None
      )
      decoder_hidden_states = (
          []
          if (return_dict_in_generate and output_hidden_states)
          else None
      )
  
      # 3. init tensors to use for "xla-compileable" generate function
      batch_size, num_beams, cur_len = shape_list(input_ids)
      # store the prompt length of decoder
      decoder_prompt_len = cur_len
  
      # per batch, beam-item holding current token in loop, pre-populated with `pad_token_id`
      input_ids_padding = tf.ones(
          (batch_size, num_beams, max_length - cur_len), dtype=tf.int32
      ) * (pad_token_id or 0)
      running_sequences = tf.concat(
          [input_ids, input_ids_padding], axis=-1
      )
      sequences = tf.ones(
          (batch_size, num_beams, max_length), dtype=tf.int32
      ) * (pad_token_id or 0)
  
      # per batch,beam-item state bit indicating if sentence has finished.
      is_sent_finished = tf.zeros(
          (batch_size, num_beams), dtype=tf.bool
      )
  
      # per batch, beam-item score, logprobs
      running_scores = tf.tile(
          tf.expand_dims(
              tf.convert_to_tensor([0.0] + [-1.0e9] * (num_beams - 1)),
              axis=0,
          ),
          [batch_size, 1],
      )
      scores = tf.ones((batch_size, num_beams)) * -1.0e9
  
      # per batch beam indices
      running_beam_indices = (
          tf.ones(
              (batch_size, num_beams, max_length - decoder_prompt_len),
              dtype=tf.int32,
          )
          * -1
      )
      beam_indices = (
          tf.ones(
              (batch_size, num_beams, max_length - decoder_prompt_len),
              dtype=tf.int32,
          )
          * -1
      )
  
      # flatten beam dim
      if "encoder_outputs" in model_kwargs:
          model_kwargs["encoder_outputs"][
              "last_hidden_state"
          ] = flatten_beam_dim(
              model_kwargs["encoder_outputs"]["last_hidden_state"]
          )
      if "attention_mask" in model_kwargs:
          model_kwargs["attention_mask"] = flatten_beam_dim(
              model_kwargs["attention_mask"]
          )
  
      # 4. define "xla-compile-able" stop-condition and auto-regressive function
      # define stop-condition and auto-regressive function
      def beam_search_cond_fn(
          cur_len,
          running_sequences,
          running_scores,
          running_beam_indices,
          sequences,
          scores,
          beam_indices,
          is_sent_finished,
          decoder_prompt_len,
          model_kwargs,
      ):
          """
          Beam Search termination condition function -- halts the generation loop if any of these conditions becomes
          False
          """
          # 1. is less than max length?
          not_max_length_yet = cur_len < max_length
  
          # 2. can the new beams still improve?
          # early_stopping == False -> apply heuristic = always get the best score from `cur_len - decoder_prompt_len`. See the discussion
          # below for more details.
          # https://github.com/huggingface/transformers/pull/20901#issuecomment-1369845565
          # early_stopping == "never" -> compute the best score from max_length or cur_len, depending on the sign of
          #   length_penalty. Positive length_penalty favors longer sequences, thus we use max_length there.
          if early_stopping == "never" and length_penalty > 0.0:
              best_running_score = running_scores[:, :1] / (
                  (max_length - decoder_prompt_len) ** length_penalty
              )
          else:
              best_running_score = running_scores[:, :1] / (
                  tf.cast(
                      cur_len - decoder_prompt_len, dtype=tf.float32
                  )
                  ** length_penalty
              )
          worst_finished_score = tf.where(
              is_sent_finished,
              tf.math.reduce_min(scores, axis=1, keepdims=True),
              -1.0e9,
          )
          improvement_still_possible = tf.math.reduce_any(
              best_running_score > worst_finished_score
          )
  
          # 3. is there still a beam that has not finished?
          still_open_beam = ~(
              tf.math.reduce_all(is_sent_finished)
              & (early_stopping is True)
          )
  
          return (
              not_max_length_yet
              & still_open_beam
              & improvement_still_possible
          )
  
      def beam_search_body_fn(
          cur_len,
          running_sequences,
          running_scores,
          running_beam_indices,
          sequences,
          scores,
          beam_indices,
          is_sent_finished,
          decoder_prompt_len,
          model_kwargs,
      ):
          """
          Beam Search iterative update function -- each iteration adds a new token and updates the best sequences
          seen so far
          """
          # 1. Forward current tokens
          if (
              model_kwargs.get("past_key_values") is None
              or needs_full_input
          ):
              input_ids = running_sequences[:, :, :cur_len]
          else:
              input_ids = tf.expand_dims(
                  running_sequences[:, :, cur_len - 1], -1
              )
          model_inputs = self.prepare_inputs_for_generation(
              flatten_beam_dim(input_ids),
              use_cache=use_cache,
              **model_kwargs,
          )
          model_outputs = self(
              **model_inputs,
              return_dict=True,
              output_attentions=output_attentions,
              output_hidden_states=output_hidden_states,
          )
          logits = unflatten_beam_dim(
              model_outputs.logits[:, -1], num_beams
          )
  
          # 2. Compute log probs
          # get log probabilities from logits, process logits with processors (*e.g.* min_length, ...), and
          # add new logprobs to existing running logprobs scores.
          log_probs = tf.nn.log_softmax(logits)
          log_probs = logits_processor(
              flatten_beam_dim(running_sequences),
              flatten_beam_dim(log_probs),
              cur_len,
          )
          log_probs = unflatten_beam_dim(log_probs, num_beams)
          if do_sample:
              log_probs = logits_warper(
                  flatten_beam_dim(running_sequences),
                  flatten_beam_dim(log_probs),
                  cur_len,
              )
              log_probs = unflatten_beam_dim(log_probs, num_beams)
          log_probs_processed = log_probs
          log_probs = log_probs + tf.expand_dims(running_scores, axis=2)
          vocab_size = log_probs.shape[2]
          log_probs = tf.reshape(
              log_probs, (batch_size, num_beams * vocab_size)
          )
  
          # Store scores, attentions and hidden_states when required
          if not use_xla and return_dict_in_generate:
              if output_scores:
                  all_scores.append(
                      logits_warper(
                          flatten_beam_dim(running_sequences),
                          flatten_beam_dim(log_probs_processed),
                          cur_len,
                      )
                  )
              if output_attentions and self.config.is_encoder_decoder:
                  decoder_attentions.append(
                      model_outputs.decoder_attentions
                  )
              elif (
                  output_attentions
                  and not self.config.is_encoder_decoder
              ):
                  decoder_attentions.append(model_outputs.attentions)
                  if self.config.is_encoder_decoder:
                      cross_attentions.append(
                          model_outputs.cross_attentions
                      )
  
              if (
                  output_hidden_states
                  and self.config.is_encoder_decoder
              ):
                  decoder_hidden_states.append(
                      model_outputs.decoder_hidden_states
                  )
              elif (
                  output_hidden_states
                  and self.config.is_encoder_decoder
              ):
                  decoder_hidden_states.append(
                      model_outputs.hidden_states
                  )
  
          # 3. Retrieve top-K
          # Each item in batch has num_beams * vocab_size candidate sequences. For each item, get the top 2*k
          # candidates with the highest log-probabilities. We gather the top 2*K beams here so that even if the
          # best K sequences reach EOS simultaneously, we have another K sequences remaining to continue the live
          # beam search.
          # Gather the top 2*K scores from _all_ beams.
          # Gather 2*k top beams.
          # Recover the beam index by floor division.
          # Recover token id by modulo division and expand Id array for broadcasting.
          # Update sequences for the 2*K top-k new sequences.
          beams_to_keep = 2 * num_beams
          if do_sample:
              topk_indices = sample_without_replacement(
                  log_probs, beams_to_keep
              )
              topk_log_probs = tf.gather(
                  log_probs, topk_indices, axis=1, batch_dims=1
              )
          else:
              topk_log_probs, topk_indices = tf.math.top_k(
                  log_probs, k=beams_to_keep
              )
          topk_current_beam_indices = topk_indices // vocab_size
          topk_running_beam_indices = self._gather_beams(
              running_beam_indices, topk_current_beam_indices
          )
          topk_running_sequences = self._gather_beams(
              running_sequences, topk_current_beam_indices
          )
          topk_ids = topk_indices % vocab_size
  
          # writes the new token
          indices_batch = tf.repeat(
              tf.range(batch_size), [beams_to_keep]
          )
          indices_beam = tf.tile(tf.range(beams_to_keep), [batch_size])
          update_indices = tf.stack(
              [
                  indices_batch,
                  indices_beam,
                  tf.broadcast_to(
                      cur_len, [batch_size * beams_to_keep]
                  ),
              ],
              axis=-1,
          )
          topk_sequences = tf.tensor_scatter_nd_update(
              tensor=topk_running_sequences,
              indices=update_indices,
              updates=tf.reshape(
                  topk_ids, [batch_size * beams_to_keep]
              ),
          )
  
          # we want to store the beam indices with batch information -> real beam index = beam index % num beams
          batch_modified_indices = (
              topk_current_beam_indices
              + tf.broadcast_to(
                  tf.expand_dims(
                      tf.range(batch_size) * num_beams, axis=1
                  ),
                  topk_current_beam_indices.shape,
              )
          )
          update_indices = tf.stack(
              [
                  indices_batch,
                  indices_beam,
                  tf.broadcast_to(
                      cur_len - decoder_prompt_len,
                      [batch_size * beams_to_keep],
                  ),
              ],
              axis=-1,
          )
          topk_beam_indices = tf.tensor_scatter_nd_update(
              tensor=topk_running_beam_indices,
              indices=update_indices,
              updates=tf.reshape(
                  batch_modified_indices, [batch_size * beams_to_keep]
              ),
          )
  
          # 4. Check which sequences have ended
          # Update current sequences: Did the top `num_beams` sequences reach an end marker?
          # To prevent these just finished sequences from being added to the current sequences
          # set of active beam search sequences, set their log probs to a very large negative value.
          if eos_token_id is None:
              eos_in_next_token = tf.zeros(
                  topk_sequences[:, :, cur_len].shape, dtype=tf.bool
              )
          else:
              eos_in_next_token = tf.math.reduce_any(
                  tf.equal(
                      tf.broadcast_to(
                          topk_sequences[:, :, cur_len],
                          [len(eos_token_id)]
                          + topk_sequences[:, :, cur_len].shape,
                      ),
                      tf.expand_dims(
                          tf.expand_dims(eos_token_id, -1), -1
                      ),
                  ),
                  axis=0,
              )
          did_topk_just_finished = eos_in_next_token & tf.broadcast_to(
              tf.concat(
                  (
                      tf.ones((num_beams), dtype=tf.bool),
                      tf.zeros((num_beams), dtype=tf.bool),
                  ),
                  axis=0,
              ),
              shape_list(eos_in_next_token),
          )
  
          # non-top `num_beams` eos tokens can't be used to finish a beam, but the others can't be used in the next
          # running sentences either
          running_topk_log_probs = (
              topk_log_probs
              + tf.cast(eos_in_next_token, tf.float32) * -1.0e9
          )
  
          # 5. Get running sequences scores for next
          # Determine the top k beam indices (from top 2*k beams) from log probs and gather top k beams
          # (from top 2*k beams).
          next_topk_indices = tf.math.top_k(
              running_topk_log_probs, k=num_beams
          )[1]
          (
              next_running_sequences,
              next_running_scores,
              next_running_beam_indices,
          ) = self._gather_beams(
              [
                  topk_sequences,
                  running_topk_log_probs,
                  topk_beam_indices,
              ],
              next_topk_indices,
          )
  
          # 6. Process topk logits
          # Further process log probs:
          # - add length penalty
          # - make sure no scores can be added anymore if beam is full
          # - make sure still running sequences cannot be chosen as finalized beam
          topk_log_probs = topk_log_probs / (
              tf.cast(
                  cur_len + 1 - decoder_prompt_len, dtype=tf.float32
              )
              ** length_penalty
          )
          beams_in_batch_are_full = tf.broadcast_to(
              tf.math.reduce_all(
                  is_sent_finished, axis=-1, keepdims=True
              ),
              shape_list(did_topk_just_finished),
          ) & (early_stopping is True)
          add_penalty = (
              ~did_topk_just_finished | beams_in_batch_are_full
          )
          topk_log_probs += tf.cast(add_penalty, tf.float32) * -1.0e9
  
          # 7. Get scores, sequences, is sentence finished for next.
          # Combine sequences, scores, and flags along the beam dimension and compare new finished sequence scores
          # to existing finished scores and select the best from the new set of beams
          merged_sequences = tf.concat(
              [sequences, topk_sequences], axis=1
          )
          merged_scores = tf.concat([scores, topk_log_probs], axis=1)
          merged_beams = tf.concat(
              [beam_indices, topk_beam_indices], axis=1
          )
          merged_is_sent_finished = tf.concat(
              [is_sent_finished, did_topk_just_finished], axis=1
          )
          topk_merged_indices = tf.math.top_k(
              merged_scores, k=num_beams
          )[1]
          (
              next_sequences,
              next_scores,
              next_beam_indices,
              next_is_sent_finished,
          ) = self._gather_beams(
              [
                  merged_sequences,
                  merged_scores,
                  merged_beams,
                  merged_is_sent_finished,
              ],
              topk_merged_indices,
          )
  
          # 8. Prepare data for the next iteration
          # Determine the top k beam indices from the original set of all beams. With these, gather the top k
          # beam-associated caches.
          cur_len = cur_len + 1
          if "past_key_values" in model_outputs:
              cache = tf.nest.map_structure(
                  lambda tensor: unflatten_beam_dim(
                      tensor, num_beams, batch_axis=cache_batch_axis
                  ),
                  model_outputs.past_key_values,
              )
              next_running_indices = self._gather_beams(
                  topk_current_beam_indices, next_topk_indices
              )
              next_cache = self._gather_beams(
                  cache,
                  next_running_indices,
                  batch_axis=cache_batch_axis,
              )
              model_outputs["past_key_values"] = tf.nest.map_structure(
                  lambda tensor: flatten_beam_dim(
                      tensor, batch_axis=cache_batch_axis
                  ),
                  next_cache,
              )
  
          if use_xla:
              next_model_kwargs = (
                  self._update_model_kwargs_for_xla_generation(
                      model_outputs=model_outputs,
                      model_kwargs=model_kwargs,
                      cur_len=cur_len,
                      max_length=max_length,
                      batch_size=(batch_size * num_beams),
                      is_encoder_decoder=self.config.is_encoder_decoder,
                      batch_axis=cache_batch_axis,
                  )
              )
          else:
              next_model_kwargs = (
                  self._update_model_kwargs_for_generation(
                      model_outputs,
                      model_kwargs,
                      is_encoder_decoder=self.config.is_encoder_decoder,
                  )
              )
  
              # if we don't cache past_key_values key values we need the whole input
              if model_kwargs.get("past_key_values", None) is None:
                  # let's throw out `past_key_values` since we don't want `None` tensors
                  model_kwargs.pop("past_key_values", None)
  
          return (
              cur_len,
              next_running_sequences,
              next_running_scores,
              next_running_beam_indices,
              next_sequences,
              next_scores,
              next_beam_indices,
              next_is_sent_finished,
              decoder_prompt_len,
              next_model_kwargs,
          )
  
      # 5. run generation
      # 1st generation step has to be run before to initialize `past_key_values` (if active)
      (
          cur_len,
          running_sequences,
          running_scores,
          running_beam_indices,
          sequences,
          scores,
          beam_indices,
          is_sent_finished,
          decoder_prompt_len,
          model_kwargs,
      ) = beam_search_body_fn(
          cur_len,
          running_sequences,
          running_scores,
          running_beam_indices,
          sequences,
          scores,
          beam_indices,
          is_sent_finished,
          decoder_prompt_len,
          model_kwargs,
      )
  
      # 2-to-n generation steps can then be run in autoregressive fashion (only in case 1st generation step does
      # NOT yield EOS token though)
      maximum_iterations = max_length - cur_len
      (
          cur_len,
          running_sequences,
          running_scores,
          running_beam_indices,
          sequences,
          scores,
          beam_indices,
          is_sent_finished,
          decoder_prompt_len,
          _,
      ) = tf.while_loop(
          beam_search_cond_fn,
          beam_search_body_fn,
          (
              cur_len,
              running_sequences,
              running_scores,
              running_beam_indices,
              sequences,
              scores,
              beam_indices,
              is_sent_finished,
              decoder_prompt_len,
              model_kwargs,
          ),
          maximum_iterations=maximum_iterations,
      )
  
      # 6. prepare outputs
      # Account for the edge-case where there are no finished sequences for a particular batch item. If so, return
      # running sequences for that batch item.
      none_finished = tf.math.reduce_any(is_sent_finished, axis=1)
      sequences = tf.where(
          none_finished[:, None, None], sequences, running_sequences
      )
      beam_indices = tf.where(
          none_finished[:, None, None],
          beam_indices,
          running_beam_indices,
      )
  
      # Apply the length penalty so that running scores match the finalized scores if they are used
      running_scores = running_scores / (
          tf.cast(cur_len - decoder_prompt_len, dtype=tf.float32)
          ** length_penalty
      )
      scores = tf.where(none_finished[:, None], scores, running_scores)
  
      # Take best beams for each batch (the score is sorted in descending order)
      sequences = flatten_beam_dim(
          sequences[:, :num_return_sequences, :]
      )
      scores = flatten_beam_dim(scores[:, :num_return_sequences])
      beam_indices = flatten_beam_dim(
          beam_indices[:, :num_return_sequences, :]
      )
  
      if not use_xla:
          # Cut for backward compatibility
          sequences = sequences[:, :cur_len]
          beam_indices = beam_indices[:, : cur_len - decoder_prompt_len]
  
      if return_dict_in_generate:
          if self.config.is_encoder_decoder:
              # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
              encoder_attentions = (
                  model_kwargs["encoder_outputs"].get("attentions")
                  if output_attentions
                  else None
              )
              encoder_hidden_states = (
                  model_kwargs["encoder_outputs"].get("hidden_states")
                  if output_hidden_states
                  else None
              )
  
              output_cls = (
                  TFBeamSampleEncoderDecoderOutput
                  if do_sample
                  else TFBeamSearchEncoderDecoderOutput
              )
              return output_cls(
                  sequences=sequences,
                  sequences_scores=scores,
                  scores=all_scores,
                  beam_indices=beam_indices,
                  encoder_attentions=encoder_attentions,
                  encoder_hidden_states=encoder_hidden_states,
                  decoder_attentions=decoder_attentions,
                  cross_attentions=cross_attentions,
                  decoder_hidden_states=decoder_hidden_states,
              )
          else:
              output_cls = (
                  TFBeamSampleDecoderOnlyOutput
                  if do_sample
                  else TFBeamSearchDecoderOnlyOutput
              )
              return output_cls(
                  sequences=sequences,
                  sequences_scores=scores,
                  scores=all_scores,
                  beam_indices=beam_indices,
                  attentions=decoder_attentions,
                  hidden_states=decoder_hidden_states,
              )
      else:
          return sequences
  
imports: import copy
  import inspect
  import warnings
  from dataclasses import dataclass
  from typing import Any, Dict, Optional, Tuple, Union
  import numpy
  import tensorflow
  from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice
  from ..modeling_tf_outputs import TFCausalLMOutputWithPast, TFSeq2SeqLMOutput
  from ..models.auto import TF_MODEL_FOR_CAUSAL_LM_MAPPING, TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, TF_MODEL_FOR_VISION_2_SEQ_MAPPING
  from ..tf_utils import shape_list, stable_softmax
  from ..utils import ModelOutput, logging
  from .configuration_utils import GenerationConfig
  from .tf_logits_process import TFForcedBOSTokenLogitsProcessor, TFForcedEOSTokenLogitsProcessor, TFForceTokensLogitsProcessor, TFLogitsProcessorList, TFMinLengthLogitsProcessor, TFNoBadWordsLogitsProcessor, TFNoRepeatNGramLogitsProcessor, TFRepetitionPenaltyLogitsProcessor, TFSuppressTokensAtBeginLogitsProcessor, TFSuppressTokensLogitsProcessor, TFTemperatureLogitsWarper, TFTopKLogitsWarper, TFTopPLogitsWarper
label_desc: Use image node_id 23 for calling the TFGenerationMixin obj's underlying member method code with example usage: obj.beam_search(input_ids, do_sample, max_length, pad_token_id, eos_token_id, length_penalty, early_stopping, logits_processor, logits_warper, num_return_sequences, output_attentions, output_hidden_states, output_scores, return_dict_in_generate) and returns: tf, tf, unknown, cur_len,  next_running_sequences,  next_running_scores,  next_running_beam_indices,  next_sequences,  next_scores,  next_beam_indices,  next_is_sent_finished,  decoder_prompt_len,  next_model_kwargs, sequences, output_cls, output_cls