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: Confirms that the model class is compatible with generation. If not, raises an exception that points to the
  right class to use.
docstr_tok: ["Confirms","that","the","model","class","is","compatible","with","generation",".","If","not",",","raises","an","exception","that","points","to","the","right","class","to","use","."]
code: def _validate_model_class(self):
      """
      Confirms that the model class is compatible with generation. If not, raises an exception that points to the
      right class to use.
      """
      if not hasattr(self, "prepare_inputs_for_generation"):
          generate_compatible_mappings = [
              FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,
              FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING,
              FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
          ]
          generate_compatible_classes = set()
          for model_mapping in generate_compatible_mappings:
              supported_models = model_mapping.get(
                  type(self.config), default=None
              )
              if supported_models is not None:
                  generate_compatible_classes.add(
                      supported_models.__name__
                  )
          exception_message = (
              f"The current model class ({self.__class__.__name__}) is not compatible with `.generate()`, as "
              "it doesn't have a language model head."
          )
          if generate_compatible_classes:
              exception_message += f" Please use one of the following classes instead: {generate_compatible_classes}"
          raise TypeError(exception_message)
  
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 5 for calling the FlaxGenerationMixin obj's underlying member method code with example usage: obj._validate_model_class() without return types