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","`","."]
code: def _greedy_search(
      self,
      input_ids: None,
      max_length: Optional[int] = None,
      pad_token_id: Optional[int] = None,
      eos_token_id: Optional[int] = None,
      logits_processor: Optional[FlaxLogitsProcessorList] = None,
      trace: bool = True,
      params: Optional[Dict[str, jnp.ndarray]] = None,
      model_kwargs: Optional[Dict[str, jnp.ndarray]] = None,
  ):
      # init values
      max_length = (
          max_length
          if max_length is not None
          else self.config.max_length
      )
      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
      )
  
      batch_size, cur_len = input_ids.shape
  
      eos_token_id = jnp.array(eos_token_id)
      pad_token_id = jnp.array(pad_token_id)
      cur_len = jnp.array(cur_len)
  
      # per batch-item holding current token in loop.
      sequences = jnp.full(
          (batch_size, max_length), pad_token_id, dtype=jnp.int32
      )
      sequences = lax.dynamic_update_slice(sequences, input_ids, (0, 0))
  
      # per batch-item state bit indicating if sentence has finished.
      is_sent_finished = jnp.zeros((batch_size,), dtype=jnp.bool_)
  
      # For Seq2Seq generation, we only need to use the decoder instead of the whole model in generation loop
      # and pass it the `encoder_outputs`, which are part of the `model_kwargs`.
      model = self.decode if self.config.is_encoder_decoder else self
      # initialize model specific kwargs
      model_kwargs = self.prepare_inputs_for_generation(
          input_ids, max_length, **model_kwargs
      )
  
      # initialize state
      state = GreedyState(
          cur_len=cur_len,
          sequences=sequences,
          running_token=input_ids,
          is_sent_finished=is_sent_finished,
          model_kwargs=model_kwargs,
      )
  
      def greedy_search_cond_fn(state):
          """state termination condition fn."""
          has_reached_max_length = state.cur_len == max_length
          all_sequence_finished = jnp.all(state.is_sent_finished)
          finish_generation = jnp.logical_or(
              has_reached_max_length, all_sequence_finished
          )
          return ~finish_generation
  
      def greedy_search_body_fn(state):
          """state update fn."""
          model_outputs = model(
              state.running_token, params=params, **state.model_kwargs
          )
          logits = model_outputs.logits[:, -1]
  
          # apply min_length, ...
          logits = logits_processor(
              state.sequences, logits, state.cur_len
          )
  
          next_token = jnp.argmax(logits, axis=-1)
  
          next_token = (
              next_token * ~state.is_sent_finished
              + pad_token_id * state.is_sent_finished
          )
          next_is_sent_finished = state.is_sent_finished | (
              next_token == eos_token_id
          )
          next_token = next_token[:, None]
  
          next_sequences = lax.dynamic_update_slice(
              state.sequences, next_token, (0, state.cur_len)
          )
          next_model_kwargs = self.update_inputs_for_generation(
              model_outputs, state.model_kwargs
          )
          return GreedyState(
              cur_len=state.cur_len + 1,
              sequences=next_sequences,
              running_token=next_token,
              is_sent_finished=next_is_sent_finished,
              model_kwargs=next_model_kwargs,
          )
  
      # The very first prompt often has sequence length > 1, so run outside of `lax.while_loop` to comply with TPU
      if input_ids.shape[1] > 1:
          state = greedy_search_body_fn(state)
  
      if not trace:
          state = self._run_loop_in_debug(
              greedy_search_cond_fn, greedy_search_body_fn, state
          )
      else:
          state = lax.while_loop(
              greedy_search_cond_fn, greedy_search_body_fn, state
          )
  
      return FlaxGreedySearchOutput(sequences=state.sequences)
  
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 10 for calling the FlaxGenerationMixin obj's underlying member method code with example usage: obj._greedy_search(input_ids, max_length, pad_token_id, eos_token_id, logits_processor, trace, params, model_kwargs) and returns: FlaxGreedySearchOutput, unknown, GreedyState