'''simple docstring'''





import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        FalconForCausalLM,
        FalconForQuestionAnswering,
        FalconForSequenceClassification,
        FalconForTokenClassification,
        FalconModel,
    )




class        __A      :



    def __init__(self    :    Dict   ,			__a    :    Optional[Any]   ,			__a    :    List[Any]=3   ,			__a    :    List[str]=7   ,			__a    :    int=True   ,			__a    :    int=True   ,			__a    :    Dict=False   ,			__a    :    List[str]=True   ,			__a    :    int=99   ,			__a    :    List[str]=32   ,			__a    :    Tuple=5   ,			__a    :    str=4   ,			__a    :    Optional[Any]=37   ,			__a    :    str="gelu"   ,			__a    :    Tuple=0.1   ,			__a    :    Optional[int]=0.1   ,			__a    :    str=512   ,			__a    :    int=16   ,			__a    :    str=2   ,			__a    :    Tuple=0.02   ,			__a    :    str=3   ,			__a    :    Tuple=4   ,			__a    :    List[str]=None   ,			):
         UpperCAmelCase_           =   parent
         UpperCAmelCase_           =   batch_size
         UpperCAmelCase_           =   seq_length
         UpperCAmelCase_           =   is_training
         UpperCAmelCase_           =   use_input_mask
         UpperCAmelCase_           =   use_token_type_ids
         UpperCAmelCase_           =   use_labels
         UpperCAmelCase_           =   vocab_size
         UpperCAmelCase_           =   hidden_size
         UpperCAmelCase_           =   num_hidden_layers
         UpperCAmelCase_           =   num_attention_heads
         UpperCAmelCase_           =   intermediate_size
         UpperCAmelCase_           =   hidden_act
         UpperCAmelCase_           =   hidden_dropout_prob
         UpperCAmelCase_           =   attention_probs_dropout_prob
         UpperCAmelCase_           =   max_position_embeddings
         UpperCAmelCase_           =   type_vocab_size
         UpperCAmelCase_           =   type_sequence_label_size
         UpperCAmelCase_           =   initializer_range
         UpperCAmelCase_           =   num_labels
         UpperCAmelCase_           =   num_choices
         UpperCAmelCase_           =   scope



    def 					_lowercase						(self    :    Union[str, Any] ):
         UpperCAmelCase_           =   ids_tensor([self.batch_size, self.seq_length]   ,			self.vocab_size )

         UpperCAmelCase_           =   None
         if self.use_input_mask:
              UpperCAmelCase_           =   random_attention_mask([self.batch_size, self.seq_length] )

         UpperCAmelCase_           =   None

         UpperCAmelCase_           =   None
         UpperCAmelCase_           =   None
         UpperCAmelCase_           =   None
         if self.use_labels:
              UpperCAmelCase_           =   ids_tensor([self.batch_size]   ,			self.type_sequence_label_size )
              UpperCAmelCase_           =   ids_tensor([self.batch_size, self.seq_length]   ,			self.num_labels )
              UpperCAmelCase_           =   ids_tensor([self.batch_size]   ,			self.num_choices )

         UpperCAmelCase_           =   self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



    def 					_lowercase						(self    :    Union[str, Any] ):
         return FalconConfig(
             vocab_size=self.vocab_size   ,			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   ,			max_position_embeddings=self.max_position_embeddings   ,			type_vocab_size=self.type_vocab_size   ,			is_decoder=__a   ,			initializer_range=self.initializer_range   ,			pad_token_id=1   ,			new_decoder_architecture=__a   ,			)



    def 					_lowercase						(self    :    Optional[int]   ,			__a    :    Union[str, Any]   ,			__a    :    int   ,			__a    :    Any   ,			__a    :    int   ,			__a    :    Union[str, Any]   ,			__a    :    str   ,			__a    :    int ):
         UpperCAmelCase_           =   FalconModel(config=__a )
         model.to(__a )
         model.eval()
         UpperCAmelCase_           =   model(__a   ,			attention_mask=__a )
         UpperCAmelCase_           =   model(__a )
         self.parent.assertEqual(result.last_hidden_state.shape   ,			(self.batch_size, self.seq_length, self.hidden_size) )



    def 					_lowercase						(self    :    Union[str, Any]   ,			__a    :    str   ,			__a    :    Optional[int]   ,			__a    :    List[Any]   ,			__a    :    Optional[Any]   ,			__a    :    str   ,			__a    :    Any   ,			__a    :    List[Any]   ,			__a    :    int   ,			__a    :    List[Any]   ,			):
         UpperCAmelCase_           =   True
         UpperCAmelCase_           =   FalconModel(__a )
         model.to(__a )
         model.eval()
         UpperCAmelCase_           =   model(
             __a   ,			attention_mask=__a   ,			encoder_hidden_states=__a   ,			encoder_attention_mask=__a   ,			)
         UpperCAmelCase_           =   model(
             __a   ,			attention_mask=__a   ,			encoder_hidden_states=__a   ,			)
         UpperCAmelCase_           =   model(__a   ,			attention_mask=__a )
         self.parent.assertEqual(result.last_hidden_state.shape   ,			(self.batch_size, self.seq_length, self.hidden_size) )



    def 					_lowercase						(self    :    Optional[Any]   ,			__a    :    Dict   ,			__a    :    Any   ,			__a    :    str   ,			__a    :    Dict   ,			__a    :    List[Any]   ,			__a    :    str   ,			__a    :    Tuple   ,			__a    :    int   ,			__a    :    List[Any]   ,			):
         UpperCAmelCase_           =   FalconForCausalLM(config=__a )
         model.to(__a )
         model.eval()
         UpperCAmelCase_           =   model(__a   ,			attention_mask=__a   ,			labels=__a )
         self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.seq_length, self.vocab_size) )



    def 					_lowercase						(self    :    Dict   ,			__a    :    int   ,			__a    :    Any   ,			__a    :    Tuple   ,			__a    :    Union[str, Any]   ,			__a    :    Optional[Any]   ,			__a    :    int   ,			__a    :    Tuple   ,			__a    :    Optional[int]   ,			__a    :    Any   ,			):
         UpperCAmelCase_           =   True
         UpperCAmelCase_           =   True
         UpperCAmelCase_           =   FalconForCausalLM(config=__a )
         model.to(__a )
         model.eval()

         # first forward pass
         UpperCAmelCase_           =   model(
             __a   ,			attention_mask=__a   ,			encoder_hidden_states=__a   ,			encoder_attention_mask=__a   ,			use_cache=__a   ,			)
         UpperCAmelCase_           =   outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         UpperCAmelCase_           =   ids_tensor((self.batch_size, 3)   ,			config.vocab_size )
         UpperCAmelCase_           =   ids_tensor((self.batch_size, 3)   ,			vocab_size=2 )

         # append to next input_ids and
         UpperCAmelCase_           =   torch.cat([input_ids, next_tokens]   ,			dim=-1 )
         UpperCAmelCase_           =   torch.cat([input_mask, next_mask]   ,			dim=-1 )

         UpperCAmelCase_           =   model(
             __a   ,			attention_mask=__a   ,			encoder_hidden_states=__a   ,			encoder_attention_mask=__a   ,			output_hidden_states=__a   ,			)["hidden_states"][0]
         UpperCAmelCase_           =   model(
             __a   ,			attention_mask=__a   ,			encoder_hidden_states=__a   ,			encoder_attention_mask=__a   ,			past_key_values=__a   ,			output_hidden_states=__a   ,			)["hidden_states"][0]

         # select random slice
         UpperCAmelCase_           =   ids_tensor((1,)   ,			output_from_past.shape[-1] ).item()
         UpperCAmelCase_           =   output_from_no_past[:, -3:, random_slice_idx].detach()
         UpperCAmelCase_           =   output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1] )

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(__a   ,			__a   ,			atol=1E-3 ) )






    def 					_lowercase						(self    :    Tuple ):
         UpperCAmelCase_           =   self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase_
         )    , (
             UpperCAmelCase_
         )    , (
             UpperCAmelCase_
         )    , (
             UpperCAmelCase_
         )    , (
             UpperCAmelCase_
         )    , (
             UpperCAmelCase_
         )    , (
             UpperCAmelCase_
         )    , 
         )           =   config_and_inputs
         UpperCAmelCase_           =   {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict




@require_torch
class        __A      ( UpperCamelCase__		,	UpperCamelCase__		,	UpperCamelCase__		,	unittest.TestCase	):
    a__       :						str          =  (
        (
            FalconModel,
            FalconForCausalLM,
            FalconForSequenceClassification,
            FalconForTokenClassification,
            FalconForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    a__       :						int          =  (FalconForCausalLM,) if is_torch_available() else ()
    a__       :						Optional[int]          =  (
        {
            """feature-extraction""": FalconModel,
            """text-classification""": FalconForSequenceClassification,
            """text-generation""": FalconForCausalLM,
            """question-answering""": FalconForQuestionAnswering,
            """token-classification""": FalconForTokenClassification,
            """zero-shot""": FalconForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    a__       :						Any          =  False
    a__       :						Tuple          =  False



    def 					_lowercase						(self    :    int ):
         UpperCAmelCase_           =   FalconModelTester(self )
         UpperCAmelCase_           =   ConfigTester(self   ,			config_class=__a   ,			hidden_size=37 )



    def 					_lowercase						(self    :    Tuple ):
         self.config_tester.run_common_tests()



    def 					_lowercase						(self    :    List[str] ):
         UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__a )



    def 					_lowercase						(self    :    Optional[int] ):
         UpperCAmelCase_    , *UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs()
         for alibi in [True, False]:
              UpperCAmelCase_           =   alibi
              self.model_tester.create_and_check_model(__a   ,			*__a )



    def 					_lowercase						(self    :    Optional[int] ):
         UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase_           =   3
         UpperCAmelCase_           =   input_dict["input_ids"]
         UpperCAmelCase_           =   input_ids.ne(1 ).to(__a )
         UpperCAmelCase_           =   ids_tensor([self.model_tester.batch_size]   ,			self.model_tester.type_sequence_label_size )
         UpperCAmelCase_           =   FalconForSequenceClassification(__a )
         model.to(__a )
         model.eval()
         UpperCAmelCase_           =   model(__a   ,			attention_mask=__a   ,			labels=__a )
         self.assertEqual(result.logits.shape   ,			(self.model_tester.batch_size, self.model_tester.num_labels) )



    def 					_lowercase						(self    :    Optional[Any] ):
         UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase_           =   3
         UpperCAmelCase_           =   "single_label_classification"
         UpperCAmelCase_           =   input_dict["input_ids"]
         UpperCAmelCase_           =   input_ids.ne(1 ).to(__a )
         UpperCAmelCase_           =   ids_tensor([self.model_tester.batch_size]   ,			self.model_tester.type_sequence_label_size )
         UpperCAmelCase_           =   FalconForSequenceClassification(__a )
         model.to(__a )
         model.eval()
         UpperCAmelCase_           =   model(__a   ,			attention_mask=__a   ,			labels=__a )
         self.assertEqual(result.logits.shape   ,			(self.model_tester.batch_size, self.model_tester.num_labels) )



    def 					_lowercase						(self    :    Any ):
         UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase_           =   input_dict["input_ids"]
         UpperCAmelCase_           =   FalconForCausalLM(__a )
         model.to(__a )
         model.eval()
         UpperCAmelCase_           =   model(__a   ,			use_cache=__a )
         UpperCAmelCase_           =   input_ids.shape[0]
         UpperCAmelCase_           =   model._convert_to_rw_cache(result.past_key_values )
         UpperCAmelCase_           =   model._convert_cache_to_standard_format(__a   ,			__a )
         for layer in range(len(__a ) ):
              for tensor_idx in range(2 ):
                   self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3 )
                   self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4 )
                   self.assertTrue(
                       torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx] ) )



    def 					_lowercase						(self    :    Optional[int] ):
         UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase_           =   3
         UpperCAmelCase_           =   "multi_label_classification"
         UpperCAmelCase_           =   input_dict["input_ids"]
         UpperCAmelCase_           =   input_ids.ne(1 ).to(__a )
         UpperCAmelCase_           =   ids_tensor(
             [self.model_tester.batch_size, config.num_labels]   ,			self.model_tester.type_sequence_label_size ).to(torch.float )
         UpperCAmelCase_           =   FalconForSequenceClassification(__a )
         model.to(__a )
         model.eval()
         UpperCAmelCase_           =   model(__a   ,			attention_mask=__a   ,			labels=__a )
         self.assertEqual(result.logits.shape   ,			(self.model_tester.batch_size, self.model_tester.num_labels) )






    def 					_lowercase						(self    :    Tuple ):
         # Falcon can have different numbers of KV-heads than the number of query heads, so we need
         # to override this test to use the right head counts.
         for model_class in self.all_generative_model_classes:
              UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()

              # If it doesn't support cache, pass the test
              if not hasattr(__a   ,			"use_cache" ):
                   return

              UpperCAmelCase_           =   model_class(__a ).to(__a )
              if "use_cache" not in inputs:
                   UpperCAmelCase_           =   True
              UpperCAmelCase_           =   model(**__a )

              # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
              if "past_key_values" not in outputs:
                   return

              UpperCAmelCase_           =   (
                  getattr(__a   ,			"decoder_layers"   ,			__a )
                  or getattr(__a   ,			"num_decoder_layers"   ,			__a )
                  or config.num_hidden_layers
              )
              UpperCAmelCase_           =   getattr(__a   ,			"num_kv_heads"   ,			config.num_attention_heads )
              UpperCAmelCase_           =   getattr(__a   ,			"d_model"   ,			config.hidden_size )
              UpperCAmelCase_           =   embed_dim // num_attention_heads

              UpperCAmelCase_           =   outputs["past_key_values"]
              self.assertEqual(len(__a )   ,			__a )

              UpperCAmelCase_    , UpperCAmelCase_           =   inputs["input_ids"].shape
              for i in range(__a ):
                   if config.new_decoder_architecture:
                        UpperCAmelCase_           =   config.num_attention_heads
                   elif config.multi_query:
                        UpperCAmelCase_           =   1
                   self.assertEqual(len(past_kv[0] )   ,			2 )  # K V for the decoder = 2
                   self.assertEqual(
                       past_kv[i][0].shape   ,			(batch_size, num_attention_heads, seq_length, per_head_embed_dim) )
                   self.assertEqual(
                       past_kv[i][1].shape   ,			(batch_size, num_attention_heads, seq_length, per_head_embed_dim) )




@require_torch
class        __A      ( unittest.TestCase	):



    @slow
    def 					_lowercase						(self    :    Union[str, Any] ):
         UpperCAmelCase_           =   AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b" )
         UpperCAmelCase_           =   FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b" )
         model.eval()
         model.to(__a )
         UpperCAmelCase_           =   tokenizer("My favorite food is"   ,			return_tensors="pt" ).to(__a )

         UpperCAmelCase_           =   (
             "My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."
         )

         UpperCAmelCase_           =   model.generate(**__a   ,			do_sample=__a   ,			max_new_tokens=19 )
         UpperCAmelCase_           =   tokenizer.batch_decode(__a )[0]

         self.assertEqual(__a   ,			__a )



    @slow
    def 					_lowercase						(self    :    Dict ):
         # The big models are way too big for the CI, so we use tiny random models that resemble their
         # architectures but with much smaller and fewer layers
         for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
              UpperCAmelCase_           =   AutoTokenizer.from_pretrained(__a )
              UpperCAmelCase_           =   FalconForCausalLM.from_pretrained(__a )
              model.eval()
              model.to(__a )
              UpperCAmelCase_           =   tokenizer("My favorite food is"   ,			return_tensors="pt" ).to(__a )

              # We just test that these run without errors - the models are randomly initialized
              # and so the actual text outputs will be garbage
              model.generate(**__a   ,			do_sample=__a   ,			max_new_tokens=4 )
              model.generate(**__a   ,			do_sample=__a   ,			max_new_tokens=4 )
              model.generate(**__a   ,			num_beams=2   ,			max_new_tokens=4 )






    @slow
    def 					_lowercase						(self    :    Union[str, Any] ):
         # The big models are way too big for the CI, so we use tiny random models that resemble their
         # architectures but with much smaller and fewer layers
         with torch.no_grad():
              for repo in [
                  "Rocketknight1/falcon-rw-1b",
                  "Rocketknight1/tiny-random-falcon-7b",
                  "Rocketknight1/tiny-random-falcon-40b",
              ]:
                   UpperCAmelCase_           =   AutoTokenizer.from_pretrained(__a )
                   UpperCAmelCase_           =   FalconForCausalLM.from_pretrained(__a )
                   model.eval()
                   model.to(device=__a )
                   UpperCAmelCase_           =   tokenizer("My favorite food is"   ,			return_tensors="pt" ).to(__a )

                   # Test results are the same with and without cache
                   UpperCAmelCase_           =   model.generate(**__a   ,			do_sample=__a   ,			max_new_tokens=20   ,			use_cache=__a )
                   UpperCAmelCase_           =   model.generate(**__a   ,			do_sample=__a   ,			max_new_tokens=20   ,			use_cache=__a )
                   self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0 )





'''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    :    str   ,			__a    :    Optional[Any]   ,			__a    :    Optional[Any]=13   ,			__a    :    int=30   ,			__a    :    Union[str, Any]=2   ,			__a    :    Dict=3   ,			__a    :    List[Any]=True   ,			__a    :    Optional[Any]=True   ,			__a    :    List[Any]=32   ,			__a    :    Any=5   ,			__a    :    str=4   ,			__a    :    Optional[int]=37   ,			__a    :    Optional[int]="gelu"   ,			__a    :    List[str]=0.1   ,			__a    :    Tuple=0.1   ,			__a    :    List[str]=10   ,			__a    :    Optional[int]=0.02   ,			):
         UpperCAmelCase_           =   parent
         UpperCAmelCase_           =   batch_size
         UpperCAmelCase_           =   image_size
         UpperCAmelCase_           =   patch_size
         UpperCAmelCase_           =   num_channels
         UpperCAmelCase_           =   is_training
         UpperCAmelCase_           =   use_labels
         UpperCAmelCase_           =   hidden_size
         UpperCAmelCase_           =   num_hidden_layers
         UpperCAmelCase_           =   num_attention_heads
         UpperCAmelCase_           =   intermediate_size
         UpperCAmelCase_           =   hidden_act
         UpperCAmelCase_           =   hidden_dropout_prob
         UpperCAmelCase_           =   attention_probs_dropout_prob
         UpperCAmelCase_           =   type_sequence_label_size
         UpperCAmelCase_           =   initializer_range

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



    def 					_lowercase						(self    :    Any ):
         UpperCAmelCase_           =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

         UpperCAmelCase_           =   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=__a   ,			initializer_range=self.initializer_range   ,			)

         return config, pixel_values



    def 					_lowercase						(self    :    Dict   ,			__a    :    Any   ,			__a    :    List[Any] ):
         UpperCAmelCase_           =   FlaxViTModel(config=__a )
         UpperCAmelCase_           =   model(__a )
         # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
         UpperCAmelCase_           =   (self.image_size, self.image_size)
         UpperCAmelCase_           =   (self.patch_size, self.patch_size)
         UpperCAmelCase_           =   (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 					_lowercase						(self    :    Tuple   ,			__a    :    str   ,			__a    :    Any ):
         UpperCAmelCase_           =   self.type_sequence_label_size
         UpperCAmelCase_           =   FlaxViTForImageClassification(config=__a )
         UpperCAmelCase_           =   model(__a )
         self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.type_sequence_label_size) )

         # test greyscale images
         UpperCAmelCase_           =   1
         UpperCAmelCase_           =   FlaxViTForImageClassification(__a )

         UpperCAmelCase_           =   floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
         UpperCAmelCase_           =   model(__a )






    def 					_lowercase						(self    :    Optional[Any] ):
         UpperCAmelCase_           =   self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase_
         )    , (
             UpperCAmelCase_
         )    , 
         )           =   config_and_inputs
         UpperCAmelCase_           =   {"pixel_values": pixel_values}
         return config, inputs_dict




@require_flax
class        __A      ( UpperCamelCase__		,	unittest.TestCase	):
    a__       :						Tuple          =  (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()



    def 					_lowercase						(self    :    Any ):
         UpperCAmelCase_           =   FlaxViTModelTester(self )
         UpperCAmelCase_           =   ConfigTester(self   ,			config_class=__a   ,			has_text_modality=__a   ,			hidden_size=37 )



    def 					_lowercase						(self    :    Tuple ):
         self.config_tester.run_common_tests()



    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__a )



    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__a )



    def 					_lowercase						(self    :    Tuple ):
         UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()

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

              UpperCAmelCase_           =   ["pixel_values"]
              self.assertListEqual(arg_names[:1]   ,			__a )



    def 					_lowercase						(self    :    Optional[Any] ):
         UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__ ):
                   UpperCAmelCase_           =   self._prepare_for_class(__a   ,			__a )
                   UpperCAmelCase_           =   model_class(__a )

                   @jax.jit
                   def model_jitted(__a    :    Tuple   ,			**__a    :    List[Any] ):
                        return model(pixel_values=__a   ,			**__a )

                   with self.subTest("JIT Enabled" ):
                        UpperCAmelCase_           =   model_jitted(**__a ).to_tuple()

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

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






    @slow
    def 					_lowercase						(self    :    Tuple ):
         for model_class_name in self.all_model_classes:
              UpperCAmelCase_           =   model_class_name.from_pretrained("google/vit-base-patch16-224" )
              UpperCAmelCase_           =   model(np.ones((1, 3, 224, 224) ) )
              self.assertIsNotNone(__a )




