






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



  def __init__(      self ,		snake_case__ ,		snake_case__=13 ,		snake_case__=30 ,		snake_case__=2 ,		snake_case__=3 ,		snake_case__=True ,		snake_case__=True ,		snake_case__=32 ,		snake_case__=5 ,		snake_case__=4 ,		snake_case__=37 ,		snake_case__="gelu" ,		snake_case__=0.1 ,		snake_case__=0.1 ,		snake_case__=10 ,		snake_case__=0.02 ,		)			->			Optional[Any]:



         """simple docstring"""





         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 		UpperCamelCase_				(      self				)			->			List[str]:



         """simple docstring"""





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

         return config, pixel_values



  def 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			FlaxViTModel(config=snake_case__				)
         UpperCAmelCase	    =			model(snake_case__				)
         # 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 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__				)			->			Tuple:



         """simple docstring"""





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

         # test greyscale images
         UpperCAmelCase	    =			1
         UpperCAmelCase	    =			FlaxViTForImageClassification(snake_case__				)

         UpperCAmelCase	    =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
         UpperCAmelCase	    =			model(snake_case__				)



  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





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





@require_flax
class   UpperCamelCase_ (       a_   ,					unittest.TestCase  ):
  _A       :					int				=	(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()



  def 		UpperCamelCase_				(      self				)			->			None:



         """simple docstring"""





         UpperCAmelCase	    =			FlaxViTModelTester(self				)
         UpperCAmelCase	    =			ConfigTester(self ,		config_class=snake_case__ ,		has_text_modality=snake_case__ ,		hidden_size=37				)



  def 		UpperCamelCase_				(      self				)			->			Optional[Any]:



         """simple docstring"""





         self.config_tester.run_common_tests()



  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*snake_case__				)



  def 		UpperCamelCase_				(      self				)			->			Any:



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*snake_case__				)



  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





         UpperCAmelCase    ,       UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
                UpperCAmelCase	    =			model_class(snake_case__				)
                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] ,		snake_case__				)



  def 		UpperCamelCase_				(      self				)			->			Dict:



         """simple docstring"""





         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(snake_case__ ,		snake_case__				)
                       UpperCAmelCase	    =			model_class(snake_case__				)

                       @jax.jit
                       def model_jitted(snake_case__ ,		**snake_case__				):
                              return model(pixel_values=snake_case__ ,		**snake_case__				)

                       with self.subTest("""JIT Enabled"""				):
                              UpperCAmelCase	    =			model_jitted(**snake_case__				).to_tuple()

                       with self.subTest("""JIT Disabled"""				):
                              with jax.disable_jit():
                                     UpperCAmelCase	    =			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				)			->			List[str]:



         """simple docstring"""





         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, 2_24, 2_24)				)				)
                self.assertIsNotNone(snake_case__				)








"""simple docstring"""







from . import __version__

# Backward compatibility imports, to make sure all those objects can be found in file_utils
from .utils import (
    CLOUDFRONT_DISTRIB_PREFIX,
    CONFIG_NAME,
    DISABLE_TELEMETRY,
    DUMMY_INPUTS,
    DUMMY_MASK,
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    FEATURE_EXTRACTOR_NAME,
    FLAX_WEIGHTS_NAME,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_PREFIX,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    MODEL_CARD_NAME,
    MULTIPLE_CHOICE_DUMMY_INPUTS,
    PYTORCH_PRETRAINED_BERT_CACHE,
    PYTORCH_TRANSFORMERS_CACHE,
    S3_BUCKET_PREFIX,
    SENTENCEPIECE_UNDERLINE,
    SPIECE_UNDERLINE,
    TF2_WEIGHTS_NAME,
    TF_WEIGHTS_NAME,
    TORCH_FX_REQUIRED_VERSION,
    TRANSFORMERS_CACHE,
    TRANSFORMERS_DYNAMIC_MODULE_NAME,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
    ContextManagers,
    DummyObject,
    EntryNotFoundError,
    ExplicitEnum,
    ModelOutput,
    PaddingStrategy,
    PushToHubMixin,
    RepositoryNotFoundError,
    RevisionNotFoundError,
    TensorType,
    _LazyModule,
    add_code_sample_docstrings,
    add_end_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    cached_property,
    copy_func,
    default_cache_path,
    define_sagemaker_information,
    get_cached_models,
    get_file_from_repo,
    get_full_repo_name,
    get_torch_version,
    has_file,
    http_user_agent,
    is_apex_available,
    is_bsa_available,
    is_coloredlogs_available,
    is_datasets_available,
    is_detectrona_available,
    is_faiss_available,
    is_flax_available,
    is_ftfy_available,
    is_in_notebook,
    is_ipex_available,
    is_librosa_available,
    is_offline_mode,
    is_onnx_available,
    is_pandas_available,
    is_phonemizer_available,
    is_protobuf_available,
    is_psutil_available,
    is_pyanvml_available,
    is_pyctcdecode_available,
    is_pytesseract_available,
    is_pytorch_quantization_available,
    is_rjieba_available,
    is_sagemaker_dp_enabled,
    is_sagemaker_mp_enabled,
    is_scipy_available,
    is_sentencepiece_available,
    is_seqio_available,
    is_sklearn_available,
    is_soundfile_availble,
    is_spacy_available,
    is_speech_available,
    is_tensor,
    is_tensorflow_probability_available,
    is_tfaonnx_available,
    is_tf_available,
    is_timm_available,
    is_tokenizers_available,
    is_torch_available,
    is_torch_bfaa_available,
    is_torch_cuda_available,
    is_torch_fx_available,
    is_torch_fx_proxy,
    is_torch_mps_available,
    is_torch_tfaa_available,
    is_torch_tpu_available,
    is_torchaudio_available,
    is_training_run_on_sagemaker,
    is_vision_available,
    replace_return_docstrings,
    requires_backends,
    to_numpy,
    to_py_obj,
    torch_only_method,
)
