



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,    __a,    __a=13,    __a=30,    __a=2,    __a=3,    __a=True,    __a=True,    __a=32,    __a=5,    __a=4,    __a=37,    __a="gelu",    __a=0.1,    __a=0.1,    __a=10,    __a=0.02,    ):

														'''simple docstring'''

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

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







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Tuple				     =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

														_lowerCAmelCase  :							List[Any]				     =					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      snake_case__   (     self,    __a,    __a):

														'''simple docstring'''

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

														'''simple docstring'''

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

														# test greyscale images
														_lowerCAmelCase  :							Any				     =					1
														_lowerCAmelCase  :							List[str]				     =					FlaxViTForImageClassification(__a)

														_lowerCAmelCase  :							Dict				     =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
														_lowerCAmelCase  :							str				     =					model(__a)







							def      snake_case__   (     self):

														'''simple docstring'''

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



@require_flax
class 						UpperCAmelCase_					( a     ,      unittest.TestCase):
							lowerCamelCase__         =	(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Tuple				     =					FlaxViTModelTester(self)
														_lowerCAmelCase  :							Union[str, Any]				     =					ConfigTester(self,    config_class=__a,    has_text_modality=__a,    hidden_size=37)







							def      snake_case__   (     self):

														'''simple docstring'''

														self.config_tester.run_common_tests()







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							int				     =					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*__a)







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							List[str]				     =					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_image_classification(*__a)







							def      snake_case__   (     self):

														'''simple docstring'''

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

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

																					_lowerCAmelCase  :							List[str]				     =					["pixel_values"]
																					self.assertListEqual(arg_names[:1],    __a)







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase     ,					_lowerCAmelCase  :							List[Any]				     =					self.model_tester.prepare_config_and_inputs_for_common()

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

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

																												with self.subTest("JIT Enabled"):
																																			_lowerCAmelCase  :							str				     =					model_jitted(**__a).to_tuple()

																												with self.subTest("JIT Disabled"):
																																			with jax.disable_jit():
																																										_lowerCAmelCase  :							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      snake_case__   (     self):

														'''simple docstring'''

														for model_class_name in self.all_model_classes:
																					_lowerCAmelCase  :							Optional[int]				     =					model_class_name.from_pretrained("google/vit-base-patch16-224")
																					_lowerCAmelCase  :							Union[str, Any]				     =					model(np.ones((1, 3, 224, 224)))
																					self.assertIsNotNone(__a)





import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING


_snake_case	=						logging.get_logger(__name__)



class 						UpperCAmelCase_					( a):
							lowerCamelCase__         =	'upernet'







							def __init__(     self,    __a=None,    __a=512,    __a=0.02,    __a=[1, 2, 3, 6],    __a=True,    __a=0.4,    __a=384,    __a=256,    __a=1,    __a=False,    __a=255,    **__a,    ):

														'''simple docstring'''

														super().__init__(**__a)

														if backbone_config is None:
																					logger.info("`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.")
																					_lowerCAmelCase  :							List[str]				     =					CONFIG_MAPPING["resnet"](out_features=["stage1", "stage2", "stage3", "stage4"])
														elif isinstance(__a,    __a):
																					_lowerCAmelCase  :							List[Any]				     =					backbone_config.get("model_type")
																					_lowerCAmelCase  :							Dict				     =					CONFIG_MAPPING[backbone_model_type]
																					_lowerCAmelCase  :							Optional[Any]				     =					config_class.from_dict(__a)

														_lowerCAmelCase  :							Tuple				     =					backbone_config
														_lowerCAmelCase  :							List[Any]				     =					hidden_size
														_lowerCAmelCase  :							Union[str, Any]				     =					initializer_range
														_lowerCAmelCase  :							str				     =					pool_scales
														_lowerCAmelCase  :							List[str]				     =					use_auxiliary_head
														_lowerCAmelCase  :							Dict				     =					auxiliary_loss_weight
														_lowerCAmelCase  :							Tuple				     =					auxiliary_in_channels
														_lowerCAmelCase  :							Optional[Any]				     =					auxiliary_channels
														_lowerCAmelCase  :							str				     =					auxiliary_num_convs
														_lowerCAmelCase  :							Union[str, Any]				     =					auxiliary_concat_input
														_lowerCAmelCase  :							Dict				     =					loss_ignore_index







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Any				     =					copy.deepcopy(self.__dict__)
														_lowerCAmelCase  :							List[Any]				     =					self.backbone_config.to_dict()
														_lowerCAmelCase  :							Optional[Any]				     =					self.__class__.model_type
														return output
