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",")","."]
code: def _sample(
      self,
      input_ids: None,
      max_length: Optional[int] = None,
      pad_token_id: Optional[int] = None,
      eos_token_id: Optional[int] = None,
      prng_key: Optional[jnp.ndarray] = None,
      logits_processor: Optional[FlaxLogitsProcessorList] = None,
      logits_warper: 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.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
      )
      prng_key = (
          prng_key if prng_key is not None else jax.random.PRNGKey(0)
      )
  
      batch_size, cur_len = input_ids.shape
  
      eos_token_id = jnp.array(
          eos_token_id,
          dtype=jnp.int32 if eos_token_id is not None else None,
      )
      pad_token_id = jnp.array(pad_token_id, dtype=jnp.int32)
      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 = SampleState(
          cur_len=cur_len,
          sequences=sequences,
          running_token=input_ids,
          is_sent_finished=is_sent_finished,
          prng_key=prng_key,
          model_kwargs=model_kwargs,
      )
  
      def sample_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 sample_search_body_fn(state):
          """state update fn."""
          prng_key, prng_key_next = jax.random.split(state.prng_key)
          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
          )
          # apply top_p, top_k, temperature
          logits = logits_warper(logits, logits, state.cur_len)
  
          next_token = jax.random.categorical(prng_key, 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 SampleState(
              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,
              prng_key=prng_key_next,
          )
  
      # 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 = sample_search_body_fn(state)
  
      if not trace:
          state = self._run_loop_in_debug(
              sample_search_cond_fn, sample_search_body_fn, state
          )
      else:
          state = lax.while_loop(
              sample_search_cond_fn, sample_search_body_fn, state
          )
  
      return FlaxSampleOutput(sequences=state.sequences)
  
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 15 for calling the FlaxGenerationMixin obj's underlying member method code with example usage: obj._sample(input_ids, max_length, pad_token_id, eos_token_id, prng_key, logits_processor, logits_warper, trace, params, model_kwargs) and returns: FlaxSampleOutput, unknown, SampleState