class_docstr: A class containing all functions for auto-regressive text generation, to be used as a mixin in
  [`FlaxPreTrainedModel`].
  
  The class exposes [`~generation_flax_utils.FlaxGenerationMixin.generate`], which can be used for:
          - *greedy decoding* by calling [`~generation_flax_utils.FlaxGenerationMixin._greedy_search`] if
            `num_beams=1` and `do_sample=False`.
          - *multinomial sampling* by calling [`~generation_flax_utils.FlaxGenerationMixin._sample`] if `num_beams=1`
            and `do_sample=True`.
          - *beam-search decoding* by calling [`~generation_utils.FlaxGenerationMixin._beam_search`] if `num_beams>1`
            and `do_sample=False`.
class_docstr_tok: ["A","class","containing","all","functions","for","auto-regressive","text","generation",",","to","be","used","as","a","mixin","in","[","`","FlaxPreTrainedModel","`","]",".","The","class","exposes","[","`","~generation_flax_utils.FlaxGenerationMixin.generate","`","]",",","which","can","be","used","for",":","-","*","greedy","decoding","*","by","calling","[","`","~generation_flax_utils.FlaxGenerationMixin._greedy_search","`","]","if","`","num_beams=1","`","and","`","do_sample=False","`",".","-","*","multinomial","sampling","*","by","calling","[","`","~generation_flax_utils.FlaxGenerationMixin._sample","`","]","if","`","num_beams=1","`","and","`","do_sample=True","`",".","-","*","beam-search","decoding","*","by","calling","[","`","~generation_utils.FlaxGenerationMixin._beam_search","`","]","if","`","num_beams",">","1","`","and","`","do_sample=False","`","."]
docstr: Generates sequences of token ids for models with a language modeling head. The method supports the following
  generation methods for text-decoder, text-to-text, speech-to-text, and vision-to-text models:
  
      - *greedy decoding* by calling [`~generation_flax_utils.FlaxGenerationMixin._greedy_search`] if
        `num_beams=1` and `do_sample=False`.
      - *multinomial sampling* by calling [`~generation_flax_utils.FlaxGenerationMixin._sample`] if `num_beams=1`
        and `do_sample=True`.
      - *beam-search decoding* by calling [`~generation_utils.FlaxGenerationMixin._beam_search`] if `num_beams>1`
        and `do_sample=False`.
  
  <Tip warning={true}>
  
  Apart from `inputs`, all the arguments below will default to the value of the attribute of the same name as
  defined in the model's config (`config.json`) which in turn defaults to the
  [`~modeling_utils.PretrainedConfig`] of the model.
  
  </Tip>
  
  Most of these parameters are explained in more detail in [this blog
  post](https://huggingface.co/blog/how-to-generate).
  
  Parameters:
      input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
          The sequence used as a prompt for the generation.
      max_length (`int`, *optional*, defaults to `model.config.max_length`):
          The maximum length the generated tokens can have. Corresponds to the length of the input prompt +
          `max_new_tokens`. In general, prefer the use of `max_new_tokens`, which ignores the number of tokens in
          the prompt.
      max_new_tokens (`int`, *optional*):
          The maximum numbers of tokens to generate, ignoring the number of tokens in the prompt.
      do_sample (`bool`, *optional*, defaults to `False`):
          Whether or not to use sampling ; use greedy decoding otherwise.
      temperature (`float`, *optional*, defaults to 1.0):
          The value used to module the next token probabilities.
      top_k (`int`, *optional*, defaults to 50):
          The number of highest probability vocabulary tokens to keep for top-k-filtering.
      top_p (`float`, *optional*, defaults to 1.0):
          If set to float < 1, only the most probable tokens with probabilities that add up to `top_p` or higher
          are kept for generation.
      pad_token_id (`int`, *optional*):
          The id of the *padding* token.
      bos_token_id (`int`, *optional*):
          The id of the *beginning-of-sequence* token.
      eos_token_id (`int`, *optional*):
          The id of the *end-of-sequence* token.
      num_beams (`int`, *optional*, defaults to 1):
          Number of beams for beam search. 1 means no beam search.
      decoder_start_token_id (`int`, *optional*):
          If an encoder-decoder model starts decoding with a different token than *bos*, the id of that token.
      trace (`bool`, *optional*, defaults to `True`):
          Whether to trace generation. Setting `trace=False` should only be used for debugging and will lead to a
          considerably slower runtime.
      params (`Dict[str, jnp.ndarray]`, *optional*):
          Optionally the model parameters can be passed. Can be useful for parallelized generation.
      model_kwargs:
          Additional model specific kwargs will be forwarded to the `forward` function of the model. If the model
          is an encoder-decoder model, encoder specific kwargs should not be prefixed and decoder specific kwargs
          should be prefixed with *decoder_*. Also accepts `encoder_outputs` to skip encoder part.
  
  Return:
      [`~utils.ModelOutput`].
  
  Examples:
  
  ```python
  >>> from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
  
  >>> tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
  >>> model = FlaxAutoModelForCausalLM.from_pretrained("distilgpt2")
  >>> input_context = "The dog"
  >>> # encode input context
  >>> input_ids = tokenizer(input_context, return_tensors="np").input_ids
  >>> # generate candidates using sampling
  >>> outputs = model.generate(input_ids=input_ids, max_length=20, top_k=30, do_sample=True)
  >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
  ```
docstr_tok: ["Generates","sequences","of","token","ids","for","models","with","a","language","modeling","head",".","The","method","supports","the","following","generation","methods","for","text-decoder",",","text-to-text",",","speech-to-text",",","and","vision-to-text","models",":","-","*","greedy","decoding","*","by","calling","[","`","~generation_flax_utils.FlaxGenerationMixin._greedy_search","`","]","if","`","num_beams=1","`","and","`","do_sample=False","`",".","-","*","multinomial","sampling","*","by","calling","[","`","~generation_flax_utils.FlaxGenerationMixin._sample","`","]","if","`","num_beams=1","`","and","`","do_sample=True","`",".","-","*","beam-search","decoding","*","by","calling","[","`","~generation_utils.FlaxGenerationMixin._beam_search","`","]","if","`","num_beams",">","1","`","and","`","do_sample=False","`",".","<","Tip","warning=","{","true","}",">","Apart","from","`","inputs","`",",","all","the","arguments","below","will","default","to","the","value","of","the","attribute","of","the","same","name","as","defined","in","the","model","'s","config","(","`","config.json","`",")","which","in","turn","defaults","to","the","[","`","~modeling_utils.PretrainedConfig","`","]","of","the","model",".","<","\/Tip",">","Most","of","these","parameters","are","explained","in","more","detail","in","[","this","blog","post","]","(","https",":","\/\/huggingface.co\/blog\/how-to-generate",")",".","Parameters",":","input_ids","(","`","jnp.ndarray","`","of","shape","`","(","batch_size",",","sequence_length",")","`",")",":","The","sequence","used","as","a","prompt","for","the","generation",".","max_length","(","`","int","`",",","*","optional","*",",","defaults","to","`","model.config.max_length","`",")",":","The","maximum","length","the","generated","tokens","can","have",".","Corresponds","to","the","length","of","the","input","prompt","+","`","max_new_tokens","`",".","In","general",",","prefer","the","use","of","`","max_new_tokens","`",",","which","ignores","the","number","of","tokens","in","the","prompt",".","max_new_tokens","(","`","int","`",",","*","optional","*",")",":","The","maximum","numbers","of","tokens","to","generate",",","ignoring","the","number","of","tokens","in","the","prompt",".","do_sample","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","or","not","to","use","sampling",";","use","greedy","decoding","otherwise",".","temperature","(","`","float","`",",","*","optional","*",",","defaults","to","1.0",")",":","The","value","used","to","module","the","next","token","probabilities",".","top_k","(","`","int","`",",","*","optional","*",",","defaults","to","50",")",":","The","number","of","highest","probability","vocabulary","tokens","to","keep","for","top-k-filtering",".","top_p","(","`","float","`",",","*","optional","*",",","defaults","to","1.0",")",":","If","set","to","float","<","1",",","only","the","most","probable","tokens","with","probabilities","that","add","up","to","`","top_p","`","or","higher","are","kept","for","generation",".","pad_token_id","(","`","int","`",",","*","optional","*",")",":","The","id","of","the","*","padding","*","token",".","bos_token_id","(","`","int","`",",","*","optional","*",")",":","The","id","of","the","*","beginning-of-sequence","*","token",".","eos_token_id","(","`","int","`",",","*","optional","*",")",":","The","id","of","the","*","end-of-sequence","*","token",".","num_beams","(","`","int","`",",","*","optional","*",",","defaults","to","1",")",":","Number","of","beams","for","beam","search",".","1","means","no","beam","search",".","decoder_start_token_id","(","`","int","`",",","*","optional","*",")",":","If","an","encoder-decoder","model","starts","decoding","with","a","different","token","than","*","bos","*",",","the","id","of","that","token",".","trace","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","to","trace","generation",".","Setting","`","trace=False","`","should","only","be","used","for","debugging","and","will","lead","to","a","considerably","slower","runtime",".","params","(","`","Dict","[","str",",","jnp.ndarray","]","`",",","*","optional","*",")",":","Optionally","the","model","parameters","can","be","passed",".","Can","be","useful","for","parallelized","generation",".","model_kwargs",":","Additional","model","specific","kwargs","will","be","forwarded","to","the","`","forward","`","function","of","the","model",".","If","the","model","is","an","encoder-decoder","model",",","encoder","specific","kwargs","should","not","be","prefixed","and","decoder","specific","kwargs","should","be","prefixed","with","*","decoder_","*",".","Also","accepts","`","encoder_outputs","`","to","skip","encoder","part",".","Return",":","[","`","~utils.ModelOutput","`","]",".","Examples",":","``","`","python",">",">",">","from","transformers","import","AutoTokenizer",",","FlaxAutoModelForCausalLM",">",">",">","tokenizer","=","AutoTokenizer.from_pretrained","(","``","distilgpt2","''",")",">",">",">","model","=","FlaxAutoModelForCausalLM.from_pretrained","(","``","distilgpt2","''",")",">",">",">","input_context","=","``","The","dog","''",">",">",">","#","encode","input","context",">",">",">","input_ids","=","tokenizer","(","input_context",",","return_tensors=","''","np","''",")",".input_ids",">",">",">","#","generate","candidates","using","sampling",">",">",">","outputs","=","model.generate","(","input_ids=input_ids",",","max_length=20",",","top_k=30",",","do_sample=True",")",">",">",">","tokenizer.batch_decode","(","outputs",",","skip_special_tokens=True",")","``","`"]
code: def generate(
      self,
      input_ids: jnp.ndarray,
      max_length: Optional[int] = None,
      max_new_tokens: Optional[int] = None,
      pad_token_id: Optional[int] = None,
      bos_token_id: Optional[int] = None,
      eos_token_id: Optional[int] = None,
      decoder_start_token_id: Optional[int] = None,
      do_sample: Optional[bool] = None,
      prng_key: Optional[jnp.ndarray] = None,
      top_k: Optional[int] = None,
      top_p: Optional[float] = None,
      temperature: Optional[float] = None,
      num_beams: Optional[int] = None,
      no_repeat_ngram_size: Optional[int] = None,
      min_length: Optional[int] = None,
      forced_bos_token_id: Optional[int] = None,
      forced_eos_token_id: Optional[int] = None,
      length_penalty: Optional[float] = None,
      early_stopping: Optional[bool] = None,
      trace: bool = True,
      params: Optional[Dict[str, jnp.ndarray]] = None,
      **model_kwargs,
  ):
      r"""
      Generates sequences of token ids for models with a language modeling head. The method supports the following
      generation methods for text-decoder, text-to-text, speech-to-text, and vision-to-text models:
  
          - *greedy decoding* by calling [`~generation_flax_utils.FlaxGenerationMixin._greedy_search`] if
            `num_beams=1` and `do_sample=False`.
          - *multinomial sampling* by calling [`~generation_flax_utils.FlaxGenerationMixin._sample`] if `num_beams=1`
            and `do_sample=True`.
          - *beam-search decoding* by calling [`~generation_utils.FlaxGenerationMixin._beam_search`] if `num_beams>1`
            and `do_sample=False`.
  
      <Tip warning={true}>
  
      Apart from `inputs`, all the arguments below will default to the value of the attribute of the same name as
      defined in the model's config (`config.json`) which in turn defaults to the
      [`~modeling_utils.PretrainedConfig`] of the model.
  
      </Tip>
  
      Most of these parameters are explained in more detail in [this blog
      post](https://huggingface.co/blog/how-to-generate).
  
      Parameters:
          input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
              The sequence used as a prompt for the generation.
          max_length (`int`, *optional*, defaults to `model.config.max_length`):
              The maximum length the generated tokens can have. Corresponds to the length of the input prompt +
              `max_new_tokens`. In general, prefer the use of `max_new_tokens`, which ignores the number of tokens in
              the prompt.
          max_new_tokens (`int`, *optional*):
              The maximum numbers of tokens to generate, ignoring the number of tokens in the prompt.
          do_sample (`bool`, *optional*, defaults to `False`):
              Whether or not to use sampling ; use greedy decoding otherwise.
          temperature (`float`, *optional*, defaults to 1.0):
              The value used to module the next token probabilities.
          top_k (`int`, *optional*, defaults to 50):
              The number of highest probability vocabulary tokens to keep for top-k-filtering.
          top_p (`float`, *optional*, defaults to 1.0):
              If set to float < 1, only the most probable tokens with probabilities that add up to `top_p` or higher
              are kept for generation.
          pad_token_id (`int`, *optional*):
              The id of the *padding* token.
          bos_token_id (`int`, *optional*):
              The id of the *beginning-of-sequence* token.
          eos_token_id (`int`, *optional*):
              The id of the *end-of-sequence* token.
          num_beams (`int`, *optional*, defaults to 1):
              Number of beams for beam search. 1 means no beam search.
          decoder_start_token_id (`int`, *optional*):
              If an encoder-decoder model starts decoding with a different token than *bos*, the id of that token.
          trace (`bool`, *optional*, defaults to `True`):
              Whether to trace generation. Setting `trace=False` should only be used for debugging and will lead to a
              considerably slower runtime.
          params (`Dict[str, jnp.ndarray]`, *optional*):
              Optionally the model parameters can be passed. Can be useful for parallelized generation.
          model_kwargs:
              Additional model specific kwargs will be forwarded to the `forward` function of the model. If the model
              is an encoder-decoder model, encoder specific kwargs should not be prefixed and decoder specific kwargs
              should be prefixed with *decoder_*. Also accepts `encoder_outputs` to skip encoder part.
  
      Return:
          [`~utils.ModelOutput`].
  
      Examples:
  
      ```python
      >>> from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
  
      >>> tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
      >>> model = FlaxAutoModelForCausalLM.from_pretrained("distilgpt2")
      >>> input_context = "The dog"
      >>> # encode input context
      >>> input_ids = tokenizer(input_context, return_tensors="np").input_ids
      >>> # generate candidates using sampling
      >>> outputs = model.generate(input_ids=input_ids, max_length=20, top_k=30, do_sample=True)
      >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
      ```"""
      # Validate the `.generate()` call
      self._validate_model_class()
      self._validate_model_kwargs(model_kwargs.copy())
  
      # set init values
      bos_token_id = (
          bos_token_id
          if bos_token_id is not None
          else self.config.bos_token_id
      )
      pad_token_id = (
          pad_token_id
          if pad_token_id is not None
          else self.config.pad_token_id
      )
      eos_token_id = (
          eos_token_id
          if eos_token_id is not None
          else self.config.eos_token_id
      )
      decoder_start_token_id = (
          decoder_start_token_id
          if decoder_start_token_id is not None
          else self.config.decoder_start_token_id
      )
      prng_key = (
          prng_key if prng_key is not None else jax.random.PRNGKey(0)
      )
  
      if (
          decoder_start_token_id is None
          and self.config.is_encoder_decoder
      ):
          raise ValueError(
              "`decoder_start_token_id` has to be defined for encoder-decoder generation."
          )
  
      # decoder-only models should use left-padding for generation (can't be checked with `trace=True`)
      if not self.config.is_encoder_decoder and not trace:
          if (
              pad_token_id is not None
              and jnp.sum(input_ids[:, -1] == pad_token_id) > 0
          ):
              logger.warning(
                  "A decoder-only architecture is being used, but right-padding was detected! For correct "
                  "generation results, please set `padding_side='left'` when initializing the tokenizer."
              )
  
      if self.config.is_encoder_decoder:
          # add encoder_outputs to model_kwargs
          if model_kwargs.get("encoder_outputs") is None:
              model_kwargs = (
                  self._prepare_encoder_decoder_kwargs_for_generation(
                      input_ids, params, model_kwargs
                  )
              )
          # prepare decoder_input_ids for generation
          input_ids = (
              jnp.ones((input_ids.shape[0], 1), dtype="i4")
              * decoder_start_token_id
          )
  
      # Prepare `max_length` depending on other stopping criteria.
      input_ids_seq_length = input_ids.shape[-1]
      if max_length is None and max_new_tokens is None:
          warnings.warn(
              "Neither `max_length` nor `max_new_tokens` have been set, `max_length` will default to "
              f"{self.config.max_length} (`self.config.max_length`). Controlling `max_length` via the config is "
              "deprecated and `max_length` will be removed from the config in v5 of Transformers -- we recommend "
              "using `max_new_tokens` to control the maximum length of the generation.",
              UserWarning,
          )
      elif max_length is None and max_new_tokens is not None:
          max_length = max_new_tokens + input_ids_seq_length
      elif max_length is not None and max_new_tokens is not None:
          raise ValueError(
              "Both `max_new_tokens` and `max_length` have been set but they serve the same purpose -- setting a"
              " limit to the generated output length. Remove one of those arguments. Please refer to the"
              " documentation for more information. "
              "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
          )
      # default to config if still None
      max_length = (
          max_length
          if max_length is not None
          else self.config.max_length
      )
      min_length = (
          min_length
          if min_length is not None
          else self.config.min_length
      )
  
      if min_length is not None and min_length > max_length:
          raise ValueError(
              f"Unfeasable length constraints: the minimum length ({min_length}) is larger than the maximum "
              f"length ({max_length})"
          )
      if input_ids_seq_length >= max_length:
          input_ids_string = (
              "decoder_input_ids"
              if self.config.is_encoder_decoder
              else "input_ids"
          )
          logger.warning(
              f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
              f" {max_length}. This can lead to unexpected behavior. You should consider increasing"
              "`max_new_tokens`."
          )
  
      do_sample = (
          do_sample if do_sample is not None else self.config.do_sample
      )
      num_beams = (
          num_beams if num_beams is not None else self.config.num_beams
      )
  
      if not do_sample and num_beams == 1:
          logits_processor = self._get_logits_processor(
              no_repeat_ngram_size,
              min_length,
              max_length,
              eos_token_id,
              forced_bos_token_id,
              forced_eos_token_id,
          )
          return self._greedy_search(
              input_ids,
              max_length,
              pad_token_id,
              eos_token_id,
              logits_processor=logits_processor,
              trace=trace,
              params=params,
              model_kwargs=model_kwargs,
          )
      elif do_sample and num_beams == 1:
          logits_warper = self._get_logits_warper(
              top_k=top_k, top_p=top_p, temperature=temperature
          )
          logits_processor = self._get_logits_processor(
              no_repeat_ngram_size,
              min_length,
              max_length,
              eos_token_id,
              forced_bos_token_id,
              forced_eos_token_id,
          )
          return self._sample(
              input_ids,
              max_length,
              pad_token_id,
              eos_token_id,
              prng_key,
              logits_warper=logits_warper,
              logits_processor=logits_processor,
              trace=trace,
              params=params,
              model_kwargs=model_kwargs,
          )
      elif not do_sample and num_beams > 1:
          # broadcast input_ids & encoder_outputs
          input_ids = self._expand_to_num_beams(
              input_ids, num_beams=num_beams
          )
  
          if "encoder_outputs" in model_kwargs:
              model_kwargs["encoder_outputs"][
                  "last_hidden_state"
              ] = self._expand_to_num_beams(
                  model_kwargs["encoder_outputs"]["last_hidden_state"],
                  num_beams=num_beams,
              )
  
          if "attention_mask" in model_kwargs:
              model_kwargs[
                  "attention_mask"
              ] = self._expand_to_num_beams(
                  model_kwargs["attention_mask"], num_beams=num_beams
              )
  
          logits_processor = self._get_logits_processor(
              no_repeat_ngram_size,
              min_length,
              max_length,
              eos_token_id,
              forced_bos_token_id,
              forced_eos_token_id,
          )
  
          return self._beam_search(
              input_ids,
              max_length,
              pad_token_id,
              eos_token_id,
              length_penalty=length_penalty,
              early_stopping=early_stopping,
              logits_processor=logits_processor,
              trace=trace,
              params=params,
              model_kwargs=model_kwargs,
          )
      else:
          raise NotImplementedError(
              "`Beam sampling is currently not implemented."
          )
  
imports: import inspect
  import warnings
  from functools import partial
  from typing import Any, Dict, Optional
  import numpy
  import flax
  import jax
  import jax.numpy
  from jax import lax
  from .generation_flax_logits_process import FlaxForcedBOSTokenLogitsProcessor, FlaxForcedEOSTokenLogitsProcessor, FlaxLogitsProcessorList, FlaxMinLengthLogitsProcessor, FlaxTemperatureLogitsWarper, FlaxTopKLogitsWarper, FlaxTopPLogitsWarper
  from .models.auto import FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING
  from .utils import ModelOutput, logging
label_desc: Use image node_id 7 for calling the FlaxGenerationMixin obj's underlying member method code with example usage: obj.generate(input_ids, max_length, max_new_tokens, pad_token_id, bos_token_id, eos_token_id, decoder_start_token_id, do_sample, prng_key, top_k, top_p, temperature, num_beams, no_repeat_ngram_size, min_length, forced_bos_token_id, forced_eos_token_id, length_penalty, early_stopping, trace, params) and returns: self, self, self