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


										'''simple docstring'''




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

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





				def  snake_case__				(						self					:						List[Any]					)							->List[str]:


										'''simple docstring'''




										_UpperCamelCase			:			Optional[int]   			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

										_UpperCamelCase			:			Tuple   			=						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  snake_case__				(						self					:						int       ,		lowercase__					:						Union[str, Any]       ,		lowercase__					:						List[Any]					)							->Optional[Any]:


										'''simple docstring'''




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


										'''simple docstring'''




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

										# test greyscale images
										_UpperCamelCase			:			str   			=						1
										_UpperCamelCase			:			Union[str, Any]   			=						FlaxViTForImageClassification(lowercase__					)

										_UpperCamelCase			:			str   			=						floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
										_UpperCamelCase			:			Optional[int]   			=						model(lowercase__					)







				def  snake_case__				(						self					:						Dict					)							->Any:


										'''simple docstring'''




										_UpperCamelCase			:			int   			=						self.prepare_config_and_inputs()
										(
										    (
										    _UpperCamelCase
										)			,       (
										    _UpperCamelCase
										)			,       
										)			:			Optional[int]   			=						config_and_inputs
										_UpperCamelCase			:			Optional[int]   			=						{"pixel_values": pixel_values}
										return config, inputs_dict





@require_flax
class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE   , unittest.TestCase	):







				'''simple docstring'''
				UpperCAmelCase__							=							(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()





				def  snake_case__				(						self					:						Dict					)							->None:


										'''simple docstring'''




										_UpperCamelCase			:			List[str]   			=						FlaxViTModelTester(self					)
										_UpperCamelCase			:			Optional[int]   			=						ConfigTester(self       ,		config_class=lowercase__       ,		has_text_modality=lowercase__       ,		hidden_size=37					)





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


										'''simple docstring'''




										self.config_tester.run_common_tests()





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


										'''simple docstring'''




										_UpperCamelCase			:			Dict   			=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowercase__					)





				def  snake_case__				(						self					:						Tuple					)							->Optional[int]:


										'''simple docstring'''




										_UpperCamelCase			:			Union[str, Any]   			=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*lowercase__					)





				def  snake_case__				(						self					:						Optional[Any]					)							->Dict:


										'''simple docstring'''




										_UpperCamelCase			,       _UpperCamelCase			:			Optional[Any]   			=						self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																_UpperCamelCase			:			Optional[int]   			=						model_class(lowercase__					)
																_UpperCamelCase			:			List[Any]   			=						inspect.signature(model.__call__					)
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																_UpperCamelCase			:			Union[str, Any]   			=						[*signature.parameters.keys()]

																_UpperCamelCase			:			Any   			=						["pixel_values"]
																self.assertListEqual(arg_names[:1]       ,		lowercase__					)





				def  snake_case__				(						self					:						Dict					)							->Tuple:


										'''simple docstring'''




										_UpperCamelCase			,       _UpperCamelCase			:			Dict   			=						self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__					):
																						_UpperCamelCase			:			str   			=						self._prepare_for_class(lowercase__       ,		lowercase__					)
																						_UpperCamelCase			:			str   			=						model_class(lowercase__					)

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

																						with self.subTest("JIT Enabled"					):
																												_UpperCamelCase			:			Optional[int]   			=						model_jitted(**lowercase__					).to_tuple()

																						with self.subTest("JIT Disabled"					):
																												with jax.disable_jit():
																																		_UpperCamelCase			:			Union[str, Any]   			=						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  snake_case__				(						self					:						List[Any]					)							->Optional[Any]:


										'''simple docstring'''




										for model_class_name in self.all_model_classes:
																_UpperCamelCase			:			Dict   			=						model_class_name.from_pretrained("google/vit-base-patch16-224"					)
																_UpperCamelCase			:			str   			=						model(np.ones((1, 3, 224, 224)					)					)
																self.assertIsNotNone(lowercase__					)



'''simple docstring'''





from collections import OrderedDict
from typing import List, Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCAmelCase_	:		Optional[Any]			=     logging.get_logger(__name__)

lowerCAmelCase_	:		List[Any]			=     {
    """google/efficientnet-b7""": """https://huggingface.co/google/efficientnet-b7/resolve/main/config.json""",
}




class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''
				UpperCAmelCase__							=							'''efficientnet'''





				def __init__(						self					:						Optional[int]       ,		lowercase__					:						int = 3       ,		lowercase__					:						int = 600       ,		lowercase__					:						float = 2.0       ,		lowercase__					:						float = 3.1       ,		lowercase__					:						int = 8       ,		lowercase__					:						List[int] = [3, 3, 5, 3, 5, 5, 3]       ,		lowercase__					:						List[int] = [32, 16, 24, 40, 80, 112, 192]       ,		lowercase__					:						List[int] = [16, 24, 40, 80, 112, 192, 320]       ,		lowercase__					:						List[int] = []       ,		lowercase__					:						List[int] = [1, 2, 2, 2, 1, 2, 1]       ,		lowercase__					:						List[int] = [1, 2, 2, 3, 3, 4, 1]       ,		lowercase__					:						List[int] = [1, 6, 6, 6, 6, 6, 6]       ,		lowercase__					:						float = 0.2_5       ,		lowercase__					:						str = "swish"       ,		lowercase__					:						int = 2_560       ,		lowercase__					:						str = "mean"       ,		lowercase__					:						float = 0.0_2       ,		lowercase__					:						float = 0.0_0_1       ,		lowercase__					:						float = 0.9_9       ,		lowercase__					:						float = 0.5       ,		lowercase__					:						float = 0.2       ,		**lowercase__					:						int       ,		)							->Optional[int]:


										'''simple docstring'''




										super().__init__(**lowercase__					)

										_UpperCamelCase			:			int   			=						num_channels
										_UpperCamelCase			:			Any   			=						image_size
										_UpperCamelCase			:			Union[str, Any]   			=						width_coefficient
										_UpperCamelCase			:			List[str]   			=						depth_coefficient
										_UpperCamelCase			:			List[Any]   			=						depth_divisor
										_UpperCamelCase			:			Optional[Any]   			=						kernel_sizes
										_UpperCamelCase			:			int   			=						in_channels
										_UpperCamelCase			:			Any   			=						out_channels
										_UpperCamelCase			:			Union[str, Any]   			=						depthwise_padding
										_UpperCamelCase			:			Optional[Any]   			=						strides
										_UpperCamelCase			:			Union[str, Any]   			=						num_block_repeats
										_UpperCamelCase			:			List[str]   			=						expand_ratios
										_UpperCamelCase			:			Dict   			=						squeeze_expansion_ratio
										_UpperCamelCase			:			Union[str, Any]   			=						hidden_act
										_UpperCamelCase			:			List[str]   			=						hidden_dim
										_UpperCamelCase			:			Dict   			=						pooling_type
										_UpperCamelCase			:			List[str]   			=						initializer_range
										_UpperCamelCase			:			Any   			=						batch_norm_eps
										_UpperCamelCase			:			List[str]   			=						batch_norm_momentum
										_UpperCamelCase			:			Optional[Any]   			=						dropout_rate
										_UpperCamelCase			:			Dict   			=						drop_connect_rate
										_UpperCamelCase			:			Dict   			=						sum(lowercase__					) * 4





class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''
				UpperCAmelCase__							=							version.parse('''1.11'''	)





				@property
				def  snake_case__				(						self					:						Union[str, Any]					)							->Mapping[str, Mapping[int, str]]:


										'''simple docstring'''




										return OrderedDict(
										    [
										        ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
										    ]					)







				@property
				def  snake_case__				(						self					:						Any					)							->float:


										'''simple docstring'''




										return 1e-5


