





from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
      from transformers import (
          TFAutoModel,
          TFAutoModelForCausalLM,
          TFAutoModelForMaskedLM,
          TFAutoModelForPreTraining,
          TFAutoModelForQuestionAnswering,
          TFAutoModelForSeqaSeqLM,
          TFAutoModelForSequenceClassification,
          TFAutoModelForTableQuestionAnswering,
          TFAutoModelForTokenClassification,
          TFAutoModelWithLMHead,
          TFBertForMaskedLM,
          TFBertForPreTraining,
          TFBertForQuestionAnswering,
          TFBertForSequenceClassification,
          TFBertModel,
          TFFunnelBaseModel,
          TFFunnelModel,
          TFGPTaLMHeadModel,
          TFRobertaForMaskedLM,
          TFTaForConditionalGeneration,
          TFTapasForQuestionAnswering,
      )
      from transformers.models.auto.modeling_tf_auto import (
          TF_MODEL_FOR_CAUSAL_LM_MAPPING,
          TF_MODEL_FOR_MASKED_LM_MAPPING,
          TF_MODEL_FOR_PRETRAINING_MAPPING,
          TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
          TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
          TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
          TF_MODEL_MAPPING,
      )
      from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
      from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
      from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
      from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST






class        snake_case_    (		__lowercase						):
  A_      							=   'new-model'


if is_tf_available():






      class        snake_case_    (		__lowercase						):
        A_      							=   NewModelConfig






@require_tf
class        snake_case_    (		unittest.TestCase						):






  @slow
  def 				UpperCAmelCase__						(							self			:					Tuple      )->List[str]:






   '''simple docstring'''

   __lowerCAmelCase					:	Tuple      =       """bert-base-cased"""
   __lowerCAmelCase					:	List[str]      =       AutoConfig.from_pretrained(_snake_case      )
   self.assertIsNotNone(_snake_case      )
   self.assertIsInstance(_snake_case       ,  _snake_case      )

   __lowerCAmelCase					:	Tuple      =       TFAutoModel.from_pretrained(_snake_case      )
   self.assertIsNotNone(_snake_case      )
   self.assertIsInstance(_snake_case       ,  _snake_case      )






  @slow
  def 				UpperCAmelCase__						(							self			:					List[Any]      )->str:






   '''simple docstring'''

   __lowerCAmelCase					:	Union[str, Any]      =       """bert-base-cased"""
   __lowerCAmelCase					:	List[Any]      =       AutoConfig.from_pretrained(_snake_case      )
   self.assertIsNotNone(_snake_case      )
   self.assertIsInstance(_snake_case       ,  _snake_case      )

   __lowerCAmelCase					:	List[Any]      =       TFAutoModelForPreTraining.from_pretrained(_snake_case      )
   self.assertIsNotNone(_snake_case      )
   self.assertIsInstance(_snake_case       ,  _snake_case      )






  @slow
  def 				UpperCAmelCase__						(							self			:					List[str]      )->int:






   '''simple docstring'''

   for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __lowerCAmelCase					:	Union[str, Any]      =       AutoConfig.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )

    __lowerCAmelCase					:	Union[str, Any]      =       TFAutoModelForCausalLM.from_pretrained(_snake_case      )
    __lowerCAmelCase ,					__lowerCAmelCase					:	Dict      =       TFAutoModelForCausalLM.from_pretrained(_snake_case       ,  output_loading_info=_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )






  @slow
  def 				UpperCAmelCase__						(							self			:					Union[str, Any]      )->str:






   '''simple docstring'''

   for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __lowerCAmelCase					:	List[str]      =       AutoConfig.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )

    __lowerCAmelCase					:	List[Any]      =       TFAutoModelWithLMHead.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )






  @slow
  def 				UpperCAmelCase__						(							self			:					Any      )->Tuple:






   '''simple docstring'''

   for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __lowerCAmelCase					:	Dict      =       AutoConfig.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )

    __lowerCAmelCase					:	List[str]      =       TFAutoModelForMaskedLM.from_pretrained(_snake_case      )
    __lowerCAmelCase ,					__lowerCAmelCase					:	str      =       TFAutoModelForMaskedLM.from_pretrained(_snake_case       ,  output_loading_info=_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )






  @slow
  def 				UpperCAmelCase__						(							self			:					List[str]      )->int:






   '''simple docstring'''

   for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __lowerCAmelCase					:	int      =       AutoConfig.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )

    __lowerCAmelCase					:	List[str]      =       TFAutoModelForSeqaSeqLM.from_pretrained(_snake_case      )
    __lowerCAmelCase ,					__lowerCAmelCase					:	List[Any]      =       TFAutoModelForSeqaSeqLM.from_pretrained(_snake_case       ,  output_loading_info=_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )






  @slow
  def 				UpperCAmelCase__						(							self			:					Optional[int]      )->Dict:






   '''simple docstring'''

   for model_name in ["bert-base-uncased"]:
    __lowerCAmelCase					:	Optional[Any]      =       AutoConfig.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )

    __lowerCAmelCase					:	Optional[Any]      =       TFAutoModelForSequenceClassification.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )






  @slow
  def 				UpperCAmelCase__						(							self			:					str      )->Dict:






   '''simple docstring'''

   for model_name in ["bert-base-uncased"]:
    __lowerCAmelCase					:	Dict      =       AutoConfig.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )

    __lowerCAmelCase					:	str      =       TFAutoModelForQuestionAnswering.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )






  @slow
  @require_tensorflow_probability
  def 				UpperCAmelCase__						(							self			:					Dict      )->int:






   '''simple docstring'''

   for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
    __lowerCAmelCase					:	List[str]      =       AutoConfig.from_pretrained(_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )

    __lowerCAmelCase					:	Optional[Any]      =       TFAutoModelForTableQuestionAnswering.from_pretrained(_snake_case      )
    __lowerCAmelCase ,					__lowerCAmelCase					:	Any      =       TFAutoModelForTableQuestionAnswering.from_pretrained(
        _snake_case       ,  output_loading_info=_snake_case      )
    self.assertIsNotNone(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )






  def 				UpperCAmelCase__						(							self			:					List[Any]      )->int:






   '''simple docstring'''

   __lowerCAmelCase					:	Any      =       TFAutoModelWithLMHead.from_pretrained(_snake_case      )
   self.assertIsInstance(_snake_case       ,  _snake_case      )
   self.assertEqual(model.num_parameters()       ,  14410      )
   self.assertEqual(model.num_parameters(only_trainable=_snake_case      )       ,  14410      )






  def 				UpperCAmelCase__						(							self			:					Dict      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	Tuple      =       TFAutoModelWithLMHead.from_pretrained(_snake_case      )
   self.assertIsInstance(_snake_case       ,  _snake_case      )
   self.assertEqual(model.num_parameters()       ,  14410      )
   self.assertEqual(model.num_parameters(only_trainable=_snake_case      )       ,  14410      )






  def 				UpperCAmelCase__						(							self			:					Optional[int]      )->List[Any]:






   '''simple docstring'''

   __lowerCAmelCase					:	Optional[Any]      =       TFAutoModel.from_pretrained("""sgugger/funnel-random-tiny"""      )
   self.assertIsInstance(_snake_case       ,  _snake_case      )

   __lowerCAmelCase					:	Optional[int]      =       copy.deepcopy(model.config      )
   __lowerCAmelCase					:	List[Any]      =       ["""FunnelBaseModel"""]
   __lowerCAmelCase					:	Optional[int]      =       TFAutoModel.from_config(_snake_case      )
   self.assertIsInstance(_snake_case       ,  _snake_case      )

   with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(_snake_case      )
    __lowerCAmelCase					:	Dict      =       TFAutoModel.from_pretrained(_snake_case      )
    self.assertIsInstance(_snake_case       ,  _snake_case      )






  def 				UpperCAmelCase__						(							self			:					List[str]      )->Any:






   '''simple docstring'''

   try:
    AutoConfig.register("""new-model"""       ,  _snake_case      )

    __lowerCAmelCase					:	int      =       [
        TFAutoModel,
        TFAutoModelForCausalLM,
        TFAutoModelForMaskedLM,
        TFAutoModelForPreTraining,
        TFAutoModelForQuestionAnswering,
        TFAutoModelForSequenceClassification,
        TFAutoModelForTokenClassification,
    ]

    for auto_class in auto_classes:
     with self.subTest(auto_class.__name__      ):
      # Wrong config class will raise an error
      with self.assertRaises(_snake_case      ):
       auto_class.register(_snake_case       ,  _snake_case      )
      auto_class.register(_snake_case       ,  _snake_case      )
      # Trying to register something existing in the Transformers library will raise an error
      with self.assertRaises(_snake_case      ):
       auto_class.register(_snake_case       ,  _snake_case      )

      # Now that the config is registered, it can be used as any other config with the auto-API
      __lowerCAmelCase					:	Union[str, Any]      =       BertModelTester(self      ).get_config()
      __lowerCAmelCase					:	Any      =       NewModelConfig(**tiny_config.to_dict()      )
      __lowerCAmelCase					:	str      =       auto_class.from_config(_snake_case      )
      self.assertIsInstance(_snake_case       ,  _snake_case      )

      with tempfile.TemporaryDirectory() as tmp_dir:
       model.save_pretrained(_snake_case      )
       __lowerCAmelCase					:	Union[str, Any]      =       auto_class.from_pretrained(_snake_case      )
       self.assertIsInstance(_snake_case       ,  _snake_case      )

   finally:
    if "new-model" in CONFIG_MAPPING._extra_content:
     del CONFIG_MAPPING._extra_content["new-model"]
    for mapping in (
        TF_MODEL_MAPPING,
        TF_MODEL_FOR_PRETRAINING_MAPPING,
        TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
        TF_MODEL_FOR_CAUSAL_LM_MAPPING,
        TF_MODEL_FOR_MASKED_LM_MAPPING,
    ):
     if NewModelConfig in mapping._extra_content:
      del mapping._extra_content[NewModelConfig]






  def 				UpperCAmelCase__						(							self			:					Any      )->Dict:






   '''simple docstring'''

   with self.assertRaisesRegex(
       _snake_case       ,  """bert-base is not a local folder and is not a valid model identifier"""      ):
    __lowerCAmelCase					:	Any      =       TFAutoModel.from_pretrained("""bert-base"""      )






  def 				UpperCAmelCase__						(							self			:					Optional[int]      )->str:






   '''simple docstring'''

   with self.assertRaisesRegex(
       _snake_case       ,  R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""      ):
    __lowerCAmelCase					:	Tuple      =       TFAutoModel.from_pretrained(_snake_case       ,  revision="""aaaaaa"""      )






  def 				UpperCAmelCase__						(							self			:					Dict      )->Any:






   '''simple docstring'''

   with self.assertRaisesRegex(
       _snake_case       ,  """hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"""       ,  ):
    __lowerCAmelCase					:	Any      =       TFAutoModel.from_pretrained("""hf-internal-testing/config-no-model"""      )






  def 				UpperCAmelCase__						(							self			:					Optional[int]      )->Union[str, Any]:






   '''simple docstring'''

   with self.assertRaisesRegex(_snake_case       ,  """Use `from_pt=True` to load this model"""      ):
    __lowerCAmelCase					:	Optional[int]      =       TFAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only"""      )






  def 				UpperCAmelCase__						(							self			:					int      )->Optional[Any]:






   '''simple docstring'''

   __lowerCAmelCase					:	str      =       TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""      )
   with RequestCounter() as counter:
    __lowerCAmelCase					:	Optional[int]      =       TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""      )
    self.assertEqual(counter.get_request_count       ,  0      )
    self.assertEqual(counter.head_request_count       ,  1      )
    self.assertEqual(counter.other_request_count       ,  0      )

   # With a sharded checkpoint
   __lowerCAmelCase					:	Dict      =       TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""      )
   with RequestCounter() as counter:
    __lowerCAmelCase					:	List[Any]      =       TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""      )
    self.assertEqual(counter.get_request_count       ,  0      )
    self.assertEqual(counter.head_request_count       ,  1      )
    self.assertEqual(counter.other_request_count       ,  0      )






import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
      import jax

      from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel






class        snake_case_    (		unittest.TestCase						):






  def __init__(							self			:					List[Any]       ,  _snake_case			:					List[Any]       ,  _snake_case			:					str=13       ,  _snake_case			:					int=30       ,  _snake_case			:					str=2       ,  _snake_case			:					int=3       ,  _snake_case			:					Optional[Any]=True       ,  _snake_case			:					str=True       ,  _snake_case			:					Optional[int]=32       ,  _snake_case			:					Dict=5       ,  _snake_case			:					Optional[int]=4       ,  _snake_case			:					List[Any]=37       ,  _snake_case			:					Union[str, Any]="gelu"       ,  _snake_case			:					str=0.1       ,  _snake_case			:					str=0.1       ,  _snake_case			:					str=10       ,  _snake_case			:					Any=0.02       ,  )->Tuple:






   '''simple docstring'''

   __lowerCAmelCase					:	Optional[int]      =       parent
   __lowerCAmelCase					:	Any      =       batch_size
   __lowerCAmelCase					:	int      =       image_size
   __lowerCAmelCase					:	int      =       patch_size
   __lowerCAmelCase					:	List[Any]      =       num_channels
   __lowerCAmelCase					:	str      =       is_training
   __lowerCAmelCase					:	str      =       use_labels
   __lowerCAmelCase					:	List[str]      =       hidden_size
   __lowerCAmelCase					:	Dict      =       num_hidden_layers
   __lowerCAmelCase					:	List[str]      =       num_attention_heads
   __lowerCAmelCase					:	Any      =       intermediate_size
   __lowerCAmelCase					:	List[str]      =       hidden_act
   __lowerCAmelCase					:	List[str]      =       hidden_dropout_prob
   __lowerCAmelCase					:	Any      =       attention_probs_dropout_prob
   __lowerCAmelCase					:	Tuple      =       type_sequence_label_size
   __lowerCAmelCase					:	Union[str, Any]      =       initializer_range

   # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
   __lowerCAmelCase					:	Tuple      =       (image_size // patch_size) ** 2
   __lowerCAmelCase					:	Optional[Any]      =       num_patches + 1






  def 				UpperCAmelCase__						(							self			:					Any      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	List[Any]      =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

   __lowerCAmelCase					:	Tuple      =       ViTConfig(
       image_size=self.image_size       ,  patch_size=self.patch_size       ,  num_channels=self.num_channels       ,  hidden_size=self.hidden_size       ,  num_hidden_layers=self.num_hidden_layers       ,  num_attention_heads=self.num_attention_heads       ,  intermediate_size=self.intermediate_size       ,  hidden_act=self.hidden_act       ,  hidden_dropout_prob=self.hidden_dropout_prob       ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,  is_decoder=_snake_case       ,  initializer_range=self.initializer_range       ,  )

   return config, pixel_values






  def 				UpperCAmelCase__						(							self			:					List[Any]       ,  _snake_case			:					List[str]       ,  _snake_case			:					List[Any]      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	str      =       FlaxViTModel(config=_snake_case      )
   __lowerCAmelCase					:	int      =       model(_snake_case      )
   # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
   __lowerCAmelCase					:	Dict      =       (self.image_size, self.image_size)
   __lowerCAmelCase					:	Any      =       (self.patch_size, self.patch_size)
   __lowerCAmelCase					:	Any      =       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
   self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, num_patches + 1, self.hidden_size)      )






  def 				UpperCAmelCase__						(							self			:					Dict       ,  _snake_case			:					str       ,  _snake_case			:					List[Any]      )->Union[str, Any]:






   '''simple docstring'''

   __lowerCAmelCase					:	Tuple      =       self.type_sequence_label_size
   __lowerCAmelCase					:	Tuple      =       FlaxViTForImageClassification(config=_snake_case      )
   __lowerCAmelCase					:	List[str]      =       model(_snake_case      )
   self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.type_sequence_label_size)      )

   # test greyscale images
   __lowerCAmelCase					:	str      =       1
   __lowerCAmelCase					:	Any      =       FlaxViTForImageClassification(_snake_case      )

   __lowerCAmelCase					:	int      =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
   __lowerCAmelCase					:	Dict      =       model(_snake_case      )






  def 				UpperCAmelCase__						(							self			:					str      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	Any      =       self.prepare_config_and_inputs()
   (
       (
       __lowerCAmelCase
   ) ,					(
       __lowerCAmelCase
   ) ,					
   )					:	Union[str, Any]      =       config_and_inputs
   __lowerCAmelCase					:	str      =       {"""pixel_values""": pixel_values}
   return config, inputs_dict






@require_flax
class        snake_case_    (		__lowercase					,unittest.TestCase						):
  A_      							=   (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()






  def 				UpperCAmelCase__						(							self			:					str      )->None:






   '''simple docstring'''

   __lowerCAmelCase					:	List[str]      =       FlaxViTModelTester(self      )
   __lowerCAmelCase					:	Optional[Any]      =       ConfigTester(self       ,  config_class=_snake_case       ,  has_text_modality=_snake_case       ,  hidden_size=37      )






  def 				UpperCAmelCase__						(							self			:					Any      )->Dict:






   '''simple docstring'''

   self.config_tester.run_common_tests()






  def 				UpperCAmelCase__						(							self			:					Any      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	Dict      =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*_snake_case      )






  def 				UpperCAmelCase__						(							self			:					Union[str, Any]      )->str:






   '''simple docstring'''

   __lowerCAmelCase					:	Dict      =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*_snake_case      )






  def 				UpperCAmelCase__						(							self			:					int      )->int:






   '''simple docstring'''

   __lowerCAmelCase ,					__lowerCAmelCase					:	Any      =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    __lowerCAmelCase					:	int      =       model_class(_snake_case      )
    __lowerCAmelCase					:	Optional[int]      =       inspect.signature(model.__call__      )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    __lowerCAmelCase					:	Tuple      =       [*signature.parameters.keys()]

    __lowerCAmelCase					:	str      =       ["""pixel_values"""]
    self.assertListEqual(arg_names[:1]       ,  _snake_case      )






  def 				UpperCAmelCase__						(							self			:					str      )->str:






   '''simple docstring'''

   __lowerCAmelCase ,					__lowerCAmelCase					:	Any      =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    with self.subTest(model_class.__name__      ):
     __lowerCAmelCase					:	Optional[int]      =       self._prepare_for_class(_snake_case       ,  _snake_case      )
     __lowerCAmelCase					:	List[str]      =       model_class(_snake_case      )

     @jax.jit
     def model_jitted(_snake_case			:					Dict       ,  **_snake_case			:					Union[str, Any]      ):
      return model(pixel_values=_snake_case       ,  **_snake_case      )

     with self.subTest("""JIT Enabled"""      ):
      __lowerCAmelCase					:	List[Any]      =       model_jitted(**_snake_case      ).to_tuple()

     with self.subTest("""JIT Disabled"""      ):
      with jax.disable_jit():
       __lowerCAmelCase					:	str      =       model_jitted(**_snake_case      ).to_tuple()

     self.assertEqual(len(_snake_case      )       ,  len(_snake_case      )      )
     for jitted_output, output in zip(_snake_case       ,  _snake_case      ):
      self.assertEqual(jitted_output.shape       ,  output.shape      )






  @slow
  def 				UpperCAmelCase__						(							self			:					Dict      )->str:






   '''simple docstring'''

   for model_class_name in self.all_model_classes:
    __lowerCAmelCase					:	List[str]      =       model_class_name.from_pretrained("""google/vit-base-patch16-224"""      )
    __lowerCAmelCase					:	List[str]      =       model(np.ones((1, 3, 224, 224)      )      )
    self.assertIsNotNone(_snake_case      )