





from . import (
    albert,
    align,
    altclip,
    audio_spectrogram_transformer,
    auto,
    autoformer,
    bark,
    bart,
    barthez,
    bartpho,
    beit,
    bert,
    bert_generation,
    bert_japanese,
    bertweet,
    big_bird,
    bigbird_pegasus,
    biogpt,
    bit,
    blenderbot,
    blenderbot_small,
    blip,
    blip_a,
    bloom,
    bridgetower,
    byta,
    camembert,
    canine,
    chinese_clip,
    clap,
    clip,
    clipseg,
    codegen,
    conditional_detr,
    convbert,
    convnext,
    convnextva,
    cpm,
    cpmant,
    ctrl,
    cvt,
    dataavec,
    deberta,
    deberta_va,
    decision_transformer,
    deformable_detr,
    deit,
    deprecated,
    deta,
    detr,
    dialogpt,
    dinat,
    distilbert,
    dit,
    donut,
    dpr,
    dpt,
    efficientformer,
    efficientnet,
    electra,
    encodec,
    encoder_decoder,
    ernie,
    ernie_m,
    esm,
    falcon,
    flaubert,
    flava,
    fnet,
    focalnet,
    fsmt,
    funnel,
    git,
    glpn,
    gpta,
    gpt_bigcode,
    gpt_neo,
    gpt_neox,
    gpt_neox_japanese,
    gpt_swa,
    gptj,
    gptsan_japanese,
    graphormer,
    groupvit,
    herbert,
    hubert,
    ibert,
    imagegpt,
    informer,
    instructblip,
    jukebox,
    layoutlm,
    layoutlmva,
    layoutlmva,
    layoutxlm,
    led,
    levit,
    lilt,
    llama,
    longformer,
    longta,
    luke,
    lxmert,
    mam_aaa,
    marian,
    markuplm,
    maskaformer,
    maskformer,
    mbart,
    mbartaa,
    mega,
    megatron_bert,
    megatron_gpta,
    mgp_str,
    mluke,
    mobilebert,
    mobilenet_va,
    mobilenet_va,
    mobilevit,
    mobilevitva,
    mpnet,
    mra,
    mta,
    musicgen,
    mvp,
    nat,
    nezha,
    nllb,
    nllb_moe,
    nystromformer,
    oneformer,
    open_llama,
    openai,
    opt,
    owlvit,
    pegasus,
    pegasus_x,
    perceiver,
    phobert,
    pixastruct,
    plbart,
    poolformer,
    prophetnet,
    qdqbert,
    rag,
    realm,
    reformer,
    regnet,
    rembert,
    resnet,
    roberta,
    roberta_prelayernorm,
    roc_bert,
    roformer,
    rwkv,
    sam,
    segformer,
    sew,
    sew_d,
    speech_encoder_decoder,
    speech_to_text,
    speech_to_text_a,
    speechta,
    splinter,
    squeezebert,
    swiftformer,
    swin,
    swinasr,
    swinva,
    switch_transformers,
    ta,
    table_transformer,
    tapas,
    time_series_transformer,
    timesformer,
    timm_backbone,
    transfo_xl,
    trocr,
    tvlt,
    umta,
    unispeech,
    unispeech_sat,
    upernet,
    videomae,
    vilt,
    vision_encoder_decoder,
    vision_text_dual_encoder,
    visual_bert,
    vit,
    vit_hybrid,
    vit_mae,
    vit_msn,
    vivit,
    wavaveca,
    wavaveca_conformer,
    wavaveca_phoneme,
    wavaveca_with_lm,
    wavlm,
    whisper,
    x_clip,
    xglm,
    xlm,
    xlm_prophetnet,
    xlm_roberta,
    xlm_roberta_xl,
    xlnet,
    xmod,
    yolos,
    yoso,
)












'''simple docstring'''

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 __a (   unittest.TestCase			):





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

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       parent
      __SCREAMING_SNAKE_CASE      							=       batch_size
      __SCREAMING_SNAKE_CASE      							=       image_size
      __SCREAMING_SNAKE_CASE      							=       patch_size
      __SCREAMING_SNAKE_CASE      							=       num_channels
      __SCREAMING_SNAKE_CASE      							=       is_training
      __SCREAMING_SNAKE_CASE      							=       use_labels
      __SCREAMING_SNAKE_CASE      							=       hidden_size
      __SCREAMING_SNAKE_CASE      							=       num_hidden_layers
      __SCREAMING_SNAKE_CASE      							=       num_attention_heads
      __SCREAMING_SNAKE_CASE      							=       intermediate_size
      __SCREAMING_SNAKE_CASE      							=       hidden_act
      __SCREAMING_SNAKE_CASE      							=       hidden_dropout_prob
      __SCREAMING_SNAKE_CASE      							=       attention_probs_dropout_prob
      __SCREAMING_SNAKE_CASE      							=       type_sequence_label_size
      __SCREAMING_SNAKE_CASE      							=       initializer_range

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





   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

      __SCREAMING_SNAKE_CASE      							=       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=lowerCamelCase				,initializer_range=self.initializer_range				,)

      return config, pixel_values





   def      UpperCAmelCase__     (					self				:			Tuple				,lowerCamelCase				:			int				,lowerCamelCase				:			Optional[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       FlaxViTModel(config=lowerCamelCase    )
      __SCREAMING_SNAKE_CASE      							=       model(lowerCamelCase    )
      # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
      __SCREAMING_SNAKE_CASE      							=       (self.image_size, self.image_size)
      __SCREAMING_SNAKE_CASE      							=       (self.patch_size, self.patch_size)
      __SCREAMING_SNAKE_CASE      							=       (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				:			Union[str, Any]				,lowerCamelCase				:			Optional[int]				,lowerCamelCase				:			Dict    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.type_sequence_label_size
      __SCREAMING_SNAKE_CASE      							=       FlaxViTForImageClassification(config=lowerCamelCase    )
      __SCREAMING_SNAKE_CASE      							=       model(lowerCamelCase    )
      self.parent.assertEqual(result.logits.shape				,(self.batch_size, self.type_sequence_label_size)    )

      # test greyscale images
      __SCREAMING_SNAKE_CASE      							=       1
      __SCREAMING_SNAKE_CASE      							=       FlaxViTForImageClassification(lowerCamelCase    )

      __SCREAMING_SNAKE_CASE      							=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
      __SCREAMING_SNAKE_CASE      							=       model(lowerCamelCase    )





   def      UpperCAmelCase__     (					self				:			int    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.prepare_config_and_inputs()
      (
          (
          __SCREAMING_SNAKE_CASE
      )							,     (
          __SCREAMING_SNAKE_CASE
      )							,     
      )      							=       config_and_inputs
      __SCREAMING_SNAKE_CASE      							=       {"""pixel_values""": pixel_values}
      return config, inputs_dict







@require_flax
class __a (   _snake_case,   unittest.TestCase			):
   __UpperCamelCase   : Any	      =	(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()





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

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       FlaxViTModelTester(self    )
      __SCREAMING_SNAKE_CASE      							=       ConfigTester(self				,config_class=lowerCamelCase				,has_text_modality=lowerCamelCase				,hidden_size=37    )





   def      UpperCAmelCase__     (					self				:			Tuple    ):

      '''simple docstring'''

      self.config_tester.run_common_tests()





   def      UpperCAmelCase__     (					self				:			List[str]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCamelCase    )





   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*lowerCamelCase    )





   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs_for_common()

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

         __SCREAMING_SNAKE_CASE      							=       ["""pixel_values"""]
         self.assertListEqual(arg_names[:1]				,lowerCamelCase    )





   def      UpperCAmelCase__     (					self				:			Optional[int]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__    ):
            __SCREAMING_SNAKE_CASE      							=       self._prepare_for_class(lowerCamelCase				,lowerCamelCase    )
            __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

            with self.subTest("""JIT Enabled"""    ):
               __SCREAMING_SNAKE_CASE      							=       model_jitted(**lowerCamelCase    ).to_tuple()

            with self.subTest("""JIT Disabled"""    ):
               with jax.disable_jit():
                  __SCREAMING_SNAKE_CASE      							=       model_jitted(**lowerCamelCase    ).to_tuple()

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





   @slow
   def      UpperCAmelCase__     (					self				:			Optional[Any]    ):

      '''simple docstring'''

      for model_class_name in self.all_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class_name.from_pretrained("""google/vit-base-patch16-224"""    )
         __SCREAMING_SNAKE_CASE      							=       model(np.ones((1, 3, 224, 224)    )    )
         self.assertIsNotNone(lowerCamelCase    )





