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)

        # 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), dtype=tf.int32) * -1
        beam_indices = tf.ones((batch_size, num_beams, max_length), 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,
            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`. 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**length_penalty)
            else:
                best_running_score = running_scores[:, :1] / (tf.cast(cur_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,
            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
            )
            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, 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,
                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,
            model_kwargs,
        ) = beam_search_body_fn(
            cur_len,
            running_sequences,
            running_scores,
            running_beam_indices,
            sequences,
            scores,
            beam_indices,
            is_sent_finished,
            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,
            _,
        ) = 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,
                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, 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]

        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

        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?']
        ```
dpkd/transformers/src/transformers/generation/tf_utils.py