class_docstr: A class containing all functions for auto-regressive text generation, to be used as a mixin in
  [`FlaxPreTrainedModel`].
  
  The class exposes [`~generation.FlaxGenerationMixin.generate`], which can be used for:
          - *greedy decoding* by calling [`~generation.FlaxGenerationMixin._greedy_search`] if `num_beams=1` and
            `do_sample=False`
          - *multinomial sampling* by calling [`~generation.FlaxGenerationMixin._sample`] if `num_beams=1` and
            `do_sample=True`
          - *beam-search decoding* by calling [`~generation.FlaxGenerationMixin._beam_search`] if `num_beams>1` and
            `do_sample=False`
  
  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","functions","for","auto-regressive","text","generation",",","to","be","used","as","a","mixin","in","[","`","FlaxPreTrainedModel","`","]",".","The","class","exposes","[","`","~generation.FlaxGenerationMixin.generate","`","]",",","which","can","be","used","for",":","-","*","greedy","decoding","*","by","calling","[","`","~generation.FlaxGenerationMixin._greedy_search","`","]","if","`","num_beams=1","`","and","`","do_sample=False","`","-","*","multinomial","sampling","*","by","calling","[","`","~generation.FlaxGenerationMixin._sample","`","]","if","`","num_beams=1","`","and","`","do_sample=True","`","-","*","beam-search","decoding","*","by","calling","[","`","~generation.FlaxGenerationMixin._beam_search","`","]","if","`","num_beams",">","1","`","and","`","do_sample=False","`","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 of token ids for models with a language modeling head.
  
  Parameters:
      input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
          The sequence used as a prompt for the generation.
      generation_config (`~generation.GenerationConfig`, *optional*):
          The generation configuration to be used as base parametrization for the generation call. `**kwargs`
          passed to generate matching the attributes of `generation_config` will override them. If
          `generation_config` is not provided, the default will be used, which had the following loading
          priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
          configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
          default values, whose documentation should be checked to parameterize generation.
      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.
      logits_processor (`FlaxLogitsProcessorList `, *optional*):
          Custom logits processors that complement the default logits processors built from arguments and
          generation config. If a logit processor is passed that is already created with the arguments or a
          generation config an error is thrown. This feature is intended for advanced users.
      kwargs (`Dict[str, Any]`, *optional*):
          Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that 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_*.
  
  Return:
      [`~utils.ModelOutput`].
docstr_tok: ["Generates","sequences","of","token","ids","for","models","with","a","language","modeling","head",".","Parameters",":","input_ids","(","`","jnp.ndarray","`","of","shape","`","(","batch_size",",","sequence_length",")","`",")",":","The","sequence","used","as","a","prompt","for","the","generation",".","generation_config","(","`","~generation.GenerationConfig","`",",","*","optional","*",")",":","The","generation","configuration","to","be","used","as","base","parametrization","for","the","generation","call",".","`","*","*","kwargs","`","passed","to","generate","matching","the","attributes","of","`","generation_config","`","will","override","them",".","If","`","generation_config","`","is","not","provided",",","the","default","will","be","used",",","which","had","the","following","loading","priority",":","1",")","from","the","`","generation_config.json","`","model","file",",","if","it","exists",";","2",")","from","the","model","configuration",".","Please","note","that","unspecified","parameters","will","inherit","[","`","~generation.GenerationConfig","`","]","'s","default","values",",","whose","documentation","should","be","checked","to","parameterize","generation",".","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",".","logits_processor","(","`","FlaxLogitsProcessorList","`",",","*","optional","*",")",":","Custom","logits","processors","that","complement","the","default","logits","processors","built","from","arguments","and","generation","config",".","If","a","logit","processor","is","passed","that","is","already","created","with","the","arguments","or","a","generation","config","an","error","is","thrown",".","This","feature","is","intended","for","advanced","users",".","kwargs","(","`","Dict","[","str",",","Any","]","`",",","*","optional","*",")",":","Ad","hoc","parametrization","of","`","generate_config","`","and\/or","additional","model-specific","kwargs","that","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_","*",".","Return",":","[","`","~utils.ModelOutput","`","]","."]
code: def generate(
      self,
      input_ids: jnp.ndarray,
      generation_config: Optional[GenerationConfig] = None,
      prng_key: Optional[jnp.ndarray] = None,
      trace: bool = True,
      params: Optional[Dict[str, jnp.ndarray]] = None,
      logits_processor: Optional[FlaxLogitsProcessorList] = None,
      **kwargs,
  ):
      r"""
      Generates sequences of token ids for models with a language modeling head.
  
      Parameters:
          input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
              The sequence used as a prompt for the generation.
          generation_config (`~generation.GenerationConfig`, *optional*):
              The generation configuration to be used as base parametrization for the generation call. `**kwargs`
              passed to generate matching the attributes of `generation_config` will override them. If
              `generation_config` is not provided, the default will be used, which had the following loading
              priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
              configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
              default values, whose documentation should be checked to parameterize generation.
          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.
          logits_processor (`FlaxLogitsProcessorList `, *optional*):
              Custom logits processors that complement the default logits processors built from arguments and
              generation config. If a logit processor is passed that is already created with the arguments or a
              generation config an error is thrown. This feature is intended for advanced users.
          kwargs (`Dict[str, Any]`, *optional*):
              Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that 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_*.
  
      Return:
          [`~utils.ModelOutput`].
  
      """
      # Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
      self._validate_model_class()
  
      # priority: `generation_config` argument > `model.generation_config` (the default generation config)
      if generation_config is None:
          # legacy: users may modify the model configuration to control generation. To trigger this legacy behavior,
          # two conditions must be met
          # 1) the generation config must have been created from the model config (`_from_model_config` field);
          # 2) the generation config must have seen no modification since its creation (the hash is the same).
          if (
              self.generation_config._from_model_config
              and self.generation_config._original_object_hash
              == hash(self.generation_config)
          ):
              new_generation_config = (
                  GenerationConfig.from_model_config(self.config)
              )
              if new_generation_config != self.generation_config:
                  warnings.warn(
                      "You have modified the pretrained model configuration to control generation. This is a"
                      " deprecated strategy to control generation and will be removed soon, in a future version."
                      " Please use and modify the model generation configuration (see"
                      " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )"
                  )
                  self.generation_config = new_generation_config
          generation_config = self.generation_config
  
      generation_config = copy.deepcopy(generation_config)
      model_kwargs = generation_config.update(
          **kwargs
      )  # All unused kwargs must be model kwargs
      generation_config.validate()
      self._validate_model_kwargs(model_kwargs.copy())
  
      logits_processor = (
          logits_processor
          if logits_processor is not None
          else FlaxLogitsProcessorList()
      )
  
      # set init values
      prng_key = (
          prng_key if prng_key is not None else jax.random.PRNGKey(0)
      )
  
      if (
          generation_config.pad_token_id is None
          and generation_config.eos_token_id is not None
      ):
          if model_kwargs.get("attention_mask") is None:
              logger.warning(
                  "The attention mask and the pad token id were not set. As a consequence, you may observe "
                  "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."
              )
          eos_token_id = generation_config.eos_token_id
          if isinstance(eos_token_id, list):
              eos_token_id = eos_token_id[0]
          logger.warning(
              f"Setting `pad_token_id` to `eos_token_id`:{eos_token_id} for open-end generation."
          )
          generation_config.pad_token_id = eos_token_id
  
      if (
          generation_config.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 (
              generation_config.pad_token_id is not None
              and jnp.sum(
                  input_ids[:, -1] == generation_config.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."
              )
  
      batch_size = input_ids.shape[0]
  
      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 = self._prepare_decoder_input_ids_for_generation(
              batch_size,
              decoder_start_token_id=generation_config.decoder_start_token_id,
              bos_token_id=generation_config.bos_token_id,
              model_kwargs=model_kwargs,
          )
  
      # Prepare `max_length` depending on other stopping criteria.
      input_ids_seq_length = input_ids.shape[-1]
      has_default_max_length = (
          kwargs.get("max_length") is None
          and generation_config.max_length is not None
      )
      if (
          has_default_max_length
          and generation_config.max_new_tokens is None
          and generation_config.max_length == 20
      ):
          # 20 is the default max_length of the generation config
          warnings.warn(
              f"Using the model-agnostic default `max_length` (={generation_config.max_length}) "
              "to control the generation length.  recommend setting `max_new_tokens` to control the maximum length of the generation.",
              UserWarning,
          )
      elif generation_config.max_new_tokens is not None:
          if (
              not has_default_max_length
              and generation_config.max_length is not None
          ):
              logger.warning(
                  f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
                  f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
                  "Please refer to the documentation for more information. "
                  "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
              )
          generation_config.max_length = (
              generation_config.max_new_tokens + input_ids_seq_length
          )
  
      if (
          generation_config.min_length is not None
          and generation_config.min_length
          > generation_config.max_length
      ):
          raise ValueError(
              f"Unfeasable length constraints: the minimum length ({generation_config.min_length}) is larger than"
              f" the maximum length ({generation_config.max_length})"
          )
      if input_ids_seq_length >= generation_config.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" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
              " increasing`max_new_tokens`."
          )
  
      logits_processor = self._get_logits_processor(
          generation_config=generation_config,
          input_ids_seq_length=input_ids_seq_length,
          logits_processor=logits_processor,
      )
  
      if (
          not generation_config.do_sample
          and generation_config.num_beams == 1
      ):
          return self._greedy_search(
              input_ids,
              generation_config.max_length,
              generation_config.pad_token_id,
              generation_config.eos_token_id,
              logits_processor=logits_processor,
              trace=trace,
              params=params,
              model_kwargs=model_kwargs,
          )
      elif (
          generation_config.do_sample
          and generation_config.num_beams == 1
      ):
          logits_warper = self._get_logits_warper(
              generation_config=generation_config
          )
          return self._sample(
              input_ids,
              generation_config.max_length,
              generation_config.pad_token_id,
              generation_config.eos_token_id,
              prng_key,
              logits_warper=logits_warper,
              logits_processor=logits_processor,
              trace=trace,
              params=params,
              model_kwargs=model_kwargs,
          )
      elif (
          not generation_config.do_sample
          and generation_config.num_beams > 1
      ):
          # broadcast input_ids & encoder_outputs
          input_ids = self._expand_to_num_beams(
              input_ids, num_beams=generation_config.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=generation_config.num_beams,
              )
  
          for kwarg in ["attention_mask", "decoder_attention_mask"]:
              if kwarg in model_kwargs:
                  model_kwargs[kwarg] = self._expand_to_num_beams(
                      model_kwargs[kwarg],
                      num_beams=generation_config.num_beams,
                  )
  
          return self._beam_search(
              input_ids,
              generation_config.max_length,
              generation_config.pad_token_id,
              generation_config.eos_token_id,
              length_penalty=generation_config.length_penalty,
              early_stopping=generation_config.early_stopping,
              logits_processor=logits_processor,
              trace=trace,
              params=params,
              num_return_sequences=generation_config.num_return_sequences,
              model_kwargs=model_kwargs,
          )
      else:
          raise NotImplementedError(
              "`Beam sampling is currently not implemented."
          )
  
imports: import copy
  import inspect
  import warnings
  from functools import partial
  from typing import Any, Dict, Optional, Union
  import flax
  import jax
  import jax.numpy
  import numpy
  from jax import lax
  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
  from .configuration_utils import GenerationConfig
  from .flax_logits_process import FlaxForcedBOSTokenLogitsProcessor, FlaxForcedEOSTokenLogitsProcessor, FlaxForceTokensLogitsProcessor, FlaxLogitsProcessorList, FlaxMinLengthLogitsProcessor, FlaxSuppressTokensAtBeginLogitsProcessor, FlaxSuppressTokensLogitsProcessor, FlaxTemperatureLogitsWarper, FlaxTopKLogitsWarper, FlaxTopPLogitsWarper
label_desc: Use image node_id 10 for calling the FlaxGenerationMixin obj's underlying member method code with example usage: obj.generate(input_ids, generation_config, prng_key, trace, params, logits_processor) and returns: self, self, self