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: Run generation in untraced mode. This should only be used for debugging purposes.
docstr_tok: ["Run","generation","in","untraced","mode",".","This","should","only","be","used","for","debugging","purposes","."]
code: def _run_loop_in_debug(cond_fn, body_fn, init_state):
      """
      Run generation in untraced mode. This should only be used for debugging purposes.
      """
      state = init_state
      while cond_fn(state):
          state = body_fn(state)
      return state
  
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 1 for calling the FlaxGenerationMixin obj's underlying member method code with example usage: obj._run_loop_in_debug(cond_fn, body_fn, init_state) and returns: state