





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



						"""simple docstring"""

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





										'''simple docstring'''







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

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







						def 		__UpperCAmelCase       (    self			:		int					):





										'''simple docstring'''







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

										_snake_case :			Dict  						=   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 		__UpperCAmelCase       (    self			:		Optional[Any] ,				lowerCamelCase_			:		Union[str, Any] ,				lowerCamelCase_			:		Optional[int]					):





										'''simple docstring'''







										_snake_case :			Any  						=   FlaxViTModel(config=A_					)
										_snake_case :			Dict  						=   model(A_					)
										# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
										_snake_case :			Optional[Any]  						=   (self.image_size, self.image_size)
										_snake_case :			str  						=   (self.patch_size, self.patch_size)
										_snake_case :			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 		__UpperCAmelCase       (    self			:		List[Any] ,				lowerCamelCase_			:		List[str] ,				lowerCamelCase_			:		Tuple					):





										'''simple docstring'''







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

										# test greyscale images
										_snake_case :			Tuple  						=   1
										_snake_case :			int  						=   FlaxViTForImageClassification(A_					)

										_snake_case :			Tuple  						=   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
										_snake_case :			List[str]  						=   model(A_					)







						def 		__UpperCAmelCase       (    self			:		Any					):





										'''simple docstring'''







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



@require_flax
class 			lowercase				(							_UpperCAmelCase		,			unittest.TestCase     ):



						"""simple docstring"""

						_UpperCamelCase  :	Any								=	(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()







						def 		__UpperCAmelCase       (    self			:		int					):





										'''simple docstring'''







										_snake_case :			int  						=   FlaxViTModelTester(self					)
										_snake_case :			Optional[int]  						=   ConfigTester(self ,				config_class=A_ ,				has_text_modality=A_ ,				hidden_size=37					)







						def 		__UpperCAmelCase       (    self			:		str					):





										'''simple docstring'''







										self.config_tester.run_common_tests()







						def 		__UpperCAmelCase       (    self			:		Optional[int]					):





										'''simple docstring'''







										_snake_case :			List[str]  						=   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*A_					)







						def 		__UpperCAmelCase       (    self			:		str					):





										'''simple docstring'''







										_snake_case :			Tuple  						=   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*A_					)







						def 		__UpperCAmelCase       (    self			:		List[Any]					):





										'''simple docstring'''







										_snake_case						,			_snake_case :			Union[str, Any]  						=   self.model_tester.prepare_config_and_inputs_for_common()

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

														_snake_case :			Optional[Any]  						=   ['pixel_values']
														self.assertListEqual(arg_names[:1] ,				A_					)







						def 		__UpperCAmelCase       (    self			:		Optional[int]					):





										'''simple docstring'''







										_snake_case						,			_snake_case :			Optional[int]  						=   self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__					):
																		_snake_case :			Union[str, Any]  						=   self._prepare_for_class(A_ ,				A_					)
																		_snake_case :			Union[str, Any]  						=   model_class(A_					)

																		@jax.jit
																		def model_jitted(lowerCamelCase_			:		str ,				**lowerCamelCase_			:		int					):
																						return model(pixel_values=A_ ,				**A_					)

																		with self.subTest('JIT Enabled'					):
																						_snake_case :			Optional[Any]  						=   model_jitted(**A_					).to_tuple()

																		with self.subTest('JIT Disabled'					):
																						with jax.disable_jit():
																										_snake_case :			Union[str, Any]  						=   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 		__UpperCAmelCase       (    self			:		int					):





										'''simple docstring'''







										for model_class_name in self.all_model_classes:
														_snake_case :			Optional[int]  						=   model_class_name.from_pretrained('google/vit-base-patch16-224'					)
														_snake_case :			Any  						=   model(np.ones((1, 3, 2_24, 2_24)					)					)
														self.assertIsNotNone(A_					)













lowercase_		:				Tuple				=						'''
# Installazione di Transformers
! pip install transformers datasets
# Per installare dalla fonte invece dell\'ultima versione rilasciata, commenta il comando sopra e
# rimuovi la modalità commento al comando seguente.
# ! pip install git+https://github.com/huggingface/transformers.git
'''

lowercase_		:				Optional[int]				=						[{'''type''': '''code''', '''content''': INSTALL_CONTENT}]
lowercase_		:				str				=						{
    '''{processor_class}''': '''FakeProcessorClass''',
    '''{model_class}''': '''FakeModelClass''',
    '''{object_class}''': '''FakeObjectClass''',
}






