


'''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      :			Dict				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Any=13				,   lowerCAmelCase      :			List[Any]=30				,   lowerCAmelCase      :			List[Any]=2				,   lowerCAmelCase      :			Any=3				,   lowerCAmelCase      :			Dict=True				,   lowerCAmelCase      :			Union[str, Any]=True				,   lowerCAmelCase      :			List[Any]=32				,   lowerCAmelCase      :			List[str]=5				,   lowerCAmelCase      :			Tuple=4				,   lowerCAmelCase      :			List[str]=37				,   lowerCAmelCase      :			str="gelu"				,   lowerCAmelCase      :			Optional[int]=0.1				,   lowerCAmelCase      :			Tuple=0.1				,   lowerCAmelCase      :			Any=10				,   lowerCAmelCase      :			Any=0.02				,   )->					Dict:

										"""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 					a__(		self      :			Dict				)->					Dict:

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

										return config, pixel_values




				def 					a__(		self      :			Dict				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Any				)->					Dict:

										"""simple docstring"""




										UpperCAmelCase					    =     FlaxViTModel(config=lowerCAmelCase				)
										UpperCAmelCase					    =     model(lowerCAmelCase				)
										# 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 					a__(		self      :			List[Any]				,   lowerCAmelCase      :			str				,   lowerCAmelCase      :			str				)->					str:

										"""simple docstring"""




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

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

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





				def 					a__(		self      :			Optional[int]				)->					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__(       lowerCAmelCase						,    unittest.TestCase	):
				__magic_name__    :				Any  		=	(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()




				def 					a__(		self      :			Union[str, Any]				)->					None:

										"""simple docstring"""




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




				def 					a__(		self      :			Dict				)->					Dict:

										"""simple docstring"""




										self.config_tester.run_common_tests()




				def 					a__(		self      :			Optional[Any]				)->					Tuple:

										"""simple docstring"""




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




				def 					a__(		self      :			Tuple				)->					Union[str, Any]:

										"""simple docstring"""




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




				def 					a__(		self      :			str				)->					Any:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()

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




				def 					a__(		self      :			Tuple				)->					Any:

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

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

																						with self.subTest('''JIT Enabled'''				):
																												UpperCAmelCase					    =     model_jitted(**lowerCAmelCase				).to_tuple()

																						with self.subTest('''JIT Disabled'''				):
																												with jax.disable_jit():
																																		UpperCAmelCase					    =     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 					a__(		self      :			Tuple				)->					Any:

										"""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, 224, 224)				)				)
																self.assertIsNotNone(lowerCAmelCase				)









'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowercase      :      List[Any]            =					{
    """configuration_x_clip""": [
        """XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """XCLIPConfig""",
        """XCLIPTextConfig""",
        """XCLIPVisionConfig""",
    ],
    """processing_x_clip""": ["""XCLIPProcessor"""],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowercase      :      List[str]            =					[
	    """XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """XCLIPModel""",
	    """XCLIPPreTrainedModel""",
	    """XCLIPTextModel""",
	    """XCLIPVisionModel""",
	]

if TYPE_CHECKING:
	from .configuration_x_clip import (
	    XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    XCLIPConfig,
	    XCLIPTextConfig,
	    XCLIPVisionConfig,
	)
	from .processing_x_clip import XCLIPProcessor

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_x_clip import (
		    XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
		    XCLIPModel,
		    XCLIPPreTrainedModel,
		    XCLIPTextModel,
		    XCLIPVisionModel,
		)

else:
	import sys

	_lowercase      :      List[Any]            =					_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





