





"""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 		SCREAMING_SNAKE_CASE_				(     unittest.TestCase ):
				"""simple docstring"""




				def __init__(  self	:Tuple    ,     __lowercase	:Union[str, Any]    ,     __lowercase	:Any=13    ,     __lowercase	:Any=30    ,     __lowercase	:Tuple=2    ,     __lowercase	:List[Any]=3    ,     __lowercase	:Any=True    ,     __lowercase	:List[Any]=True    ,     __lowercase	:List[Any]=32    ,     __lowercase	:Optional[int]=5    ,     __lowercase	:Tuple=4    ,     __lowercase	:Union[str, Any]=37    ,     __lowercase	:List[Any]="gelu"    ,     __lowercase	:Optional[int]=0.1    ,     __lowercase	:Optional[Any]=0.1    ,     __lowercase	:Any=10    ,     __lowercase	:Union[str, Any]=0.02    ,     ):
											__lowerCamelCase   :   Dict       =parent
											__lowerCamelCase   :   str       =batch_size
											__lowerCamelCase   :   Dict       =image_size
											__lowerCamelCase   :   List[str]       =patch_size
											__lowerCamelCase   :   List[Any]       =num_channels
											__lowerCamelCase   :   Union[str, Any]       =is_training
											__lowerCamelCase   :   Tuple       =use_labels
											__lowerCamelCase   :   int       =hidden_size
											__lowerCamelCase   :   Union[str, Any]       =num_hidden_layers
											__lowerCamelCase   :   Any       =num_attention_heads
											__lowerCamelCase   :   Tuple       =intermediate_size
											__lowerCamelCase   :   str       =hidden_act
											__lowerCamelCase   :   List[Any]       =hidden_dropout_prob
											__lowerCamelCase   :   Optional[int]       =attention_probs_dropout_prob
											__lowerCamelCase   :   str       =type_sequence_label_size
											__lowerCamelCase   :   int       =initializer_range

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




				def        __lowercase	(  self	:str      ):
											__lowerCamelCase   :   str       =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

											__lowerCamelCase   :   Optional[int]       =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=__lowercase    ,     initializer_range=self.initializer_range    ,     )

											return config, pixel_values




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

											# test greyscale images
											__lowerCamelCase   :   int       =1
											__lowerCamelCase   :   Optional[Any]       =FlaxViTForImageClassification(__lowercase      )

											__lowerCamelCase   :   List[str]       =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
											__lowerCamelCase   :   Dict       =model(__lowercase      )




				def        __lowercase	(  self	:Union[str, Any]      ):
											__lowerCamelCase   :   str       =self.prepare_config_and_inputs()
											(
											    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    
											)   :   Dict       =config_and_inputs
											__lowerCamelCase   :   str       ={'''pixel_values''': pixel_values}
											return config, inputs_dict

@require_flax
class 		SCREAMING_SNAKE_CASE_				(     snake_case__							,	unittest.TestCase ):
				"""simple docstring"""
				__snake_case :							Any									=		(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()




				def        __lowercase	(  self	:Optional[Any]      ):
											__lowerCamelCase   :   Tuple       =FlaxViTModelTester(self      )
											__lowerCamelCase   :   Tuple       =ConfigTester(self    ,     config_class=__lowercase    ,     has_text_modality=__lowercase    ,     hidden_size=37      )




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




				def        __lowercase	(  self	:Tuple      ):
											__lowerCamelCase   :   List[str]       =self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__lowercase      )




				def        __lowercase	(  self	:int      ):
											__lowerCamelCase   :   int       =self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*__lowercase      )




				def        __lowercase	(  self	:str      ):
											__lowerCamelCase     ,    __lowerCamelCase   :   int       =self.model_tester.prepare_config_and_inputs_for_common()

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

																		__lowerCamelCase   :   List[str]       =['''pixel_values''']
																		self.assertListEqual(arg_names[:1]    ,     __lowercase      )




				def        __lowercase	(  self	:Optional[Any]      ):
											__lowerCamelCase     ,    __lowerCamelCase   :   Tuple       =self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																		with self.subTest(model_class.__name__      ):
																									__lowerCamelCase   :   Any       =self._prepare_for_class(__lowercase    ,     __lowercase      )
																									__lowerCamelCase   :   str       =model_class(__lowercase      )

																									@jax.jit
																									def model_jitted(__lowercase	:int    ,     **__lowercase	:Optional[int]      ):
																																return model(pixel_values=__lowercase    ,     **__lowercase      )

																									with self.subTest('''JIT Enabled'''      ):
																																__lowerCamelCase   :   Dict       =model_jitted(**__lowercase      ).to_tuple()

																									with self.subTest('''JIT Disabled'''      ):
																																with jax.disable_jit():
																																							__lowerCamelCase   :   List[str]       =model_jitted(**__lowercase      ).to_tuple()

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




				@slow
				def        __lowercase	(  self	:str      ):
											for model_class_name in self.all_model_classes:
																		__lowerCamelCase   :   Any       =model_class_name.from_pretrained('''google/vit-base-patch16-224'''      )
																		__lowerCamelCase   :   Optional[Any]       =model(np.ones((1, 3, 224, 224)      )      )
																		self.assertIsNotNone(__lowercase      )











"""simple docstring"""





from packaging import version

from .import_utils import is_accelerate_available


if is_accelerate_available():
							import accelerate



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				str ):


							'''simple docstring'''







							if not is_accelerate_available():
														return method
							__lowerCamelCase   :   Dict       =version.parse(accelerate.__version__ ).base_version
							if version.parse(SCREAMING_SNAKE_CASE ) < version.parse('''0.17.0''' ):
														return method

							def wrapper(self							:				List[Any]     ,    *SCREAMING_SNAKE_CASE							:				Any     ,    **SCREAMING_SNAKE_CASE							:				List[str] ):
														if hasattr(self     ,    '''_hf_hook''' ) and hasattr(self._hf_hook     ,    '''pre_forward''' ):
																					self._hf_hook.pre_forward(self )
														return method(self     ,    *SCREAMING_SNAKE_CASE     ,    **SCREAMING_SNAKE_CASE )

							return wrapper




