






'''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 ):




    '''simple docstring'''







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

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






    def       _UpperCamelCase    (       self )							-> Any:
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

         # test greyscale images
         SCREAMING_SNAKE_CASE					:			Any    	=       1
         SCREAMING_SNAKE_CASE					:			Tuple    	=       FlaxViTForImageClassification(lowercase__ )

         SCREAMING_SNAKE_CASE					:			List[str]    	=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       model(lowercase__ )






    def       _UpperCamelCase    (       self )							-> Optional[int]:
         SCREAMING_SNAKE_CASE					:			int    	=       self.prepare_config_and_inputs()
         (
             SCREAMING_SNAKE_CASE   
         )					:			List[str]    	=       config_and_inputs
         SCREAMING_SNAKE_CASE					:			Any    	=       {'pixel_values': pixel_values}
         return config, inputs_dict






@require_flax
class 				UpperCAmelCase (       _SCREAMING_SNAKE_CASE   ,    unittest.TestCase ):




    '''simple docstring'''







    snake_case__							:   Dict							=						(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()






    def       _UpperCamelCase    (       self )							-> None:
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       FlaxViTModelTester(self )
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       ConfigTester(self    ,							config_class=lowercase__    ,							has_text_modality=lowercase__    ,							hidden_size=37 )






    def       _UpperCamelCase    (       self )							-> Any:
         self.config_tester.run_common_tests()






    def       _UpperCamelCase    (       self )							-> str:
         SCREAMING_SNAKE_CASE					:			Dict    	=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowercase__ )






    def       _UpperCamelCase    (       self )							-> List[Any]:
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*lowercase__ )






    def       _UpperCamelCase    (       self )							-> Tuple:
         SCREAMING_SNAKE_CASE					:			List[str]    	=       self.model_tester.prepare_config_and_inputs_for_common()

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

              SCREAMING_SNAKE_CASE					:			List[Any]    	=       ['pixel_values']
              self.assertListEqual(arg_names[:1]    ,							lowercase__ )






    def       _UpperCamelCase    (       self )							-> List[Any]:
         SCREAMING_SNAKE_CASE					:			int    	=       self.model_tester.prepare_config_and_inputs_for_common()

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

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

                   with self.subTest('JIT Enabled' ):
                        SCREAMING_SNAKE_CASE					:			Any    	=       model_jitted(**lowercase__ ).to_tuple()

                   with self.subTest('JIT Disabled' ):
                        with jax.disable_jit():
                             SCREAMING_SNAKE_CASE					:			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       _UpperCamelCase    (       self )							-> int:
         for model_class_name in self.all_model_classes:
              SCREAMING_SNAKE_CASE					:			str    	=       model_class_name.from_pretrained('google/vit-base-patch16-224' )
              SCREAMING_SNAKE_CASE					:			Tuple    	=       model(np.ones((1, 3, 224, 224) ) )
              self.assertIsNotNone(lowercase__ )




'''simple docstring'''





import warnings
from collections import OrderedDict
from typing import Mapping

from packaging import version

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


_lowerCAmelCase  :str  				=    logging.get_logger(__name__)

_lowerCAmelCase  :Optional[int]  				=    {
    """nvidia/segformer-b0-finetuned-ade-512-512""": (
        """https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512/resolve/main/config.json"""
    ),
    # See all SegFormer models at https://huggingface.co/models?filter=segformer
}






class 				UpperCAmelCase (       _SCREAMING_SNAKE_CASE ):




    '''simple docstring'''







    snake_case__							:   Optional[int]							=						"segformer"






    def __init__(       self    ,							lowercase__=3    ,							lowercase__=4    ,							lowercase__=[2, 2, 2, 2]    ,							lowercase__=[8, 4, 2, 1]    ,							lowercase__=[32, 64, 160, 256]    ,							lowercase__=[7, 3, 3, 3]    ,							lowercase__=[4, 2, 2, 2]    ,							lowercase__=[1, 2, 5, 8]    ,							lowercase__=[4, 4, 4, 4]    ,							lowercase__="gelu"    ,							lowercase__=0.0    ,							lowercase__=0.0    ,							lowercase__=0.1    ,							lowercase__=0.0_2    ,							lowercase__=0.1    ,							lowercase__=1E-6    ,							lowercase__=256    ,							lowercase__=255    ,							**lowercase__    ,							)							-> Union[str, Any]:
         super().__init__(**lowercase__ )

         if "reshape_last_stage" in kwargs and kwargs["reshape_last_stage"] is False:
              warnings.warn(
                  'Reshape_last_stage is set to False in this config. This argument is deprecated and will soon be'
                  ' removed, as the behaviour will default to that of reshape_last_stage = True.'    ,							lowercase__    ,							)

         SCREAMING_SNAKE_CASE					:			int    	=       num_channels
         SCREAMING_SNAKE_CASE					:			int    	=       num_encoder_blocks
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       depths
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       sr_ratios
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       hidden_sizes
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       patch_sizes
         SCREAMING_SNAKE_CASE					:			int    	=       strides
         SCREAMING_SNAKE_CASE					:			Tuple    	=       mlp_ratios
         SCREAMING_SNAKE_CASE					:			List[str]    	=       num_attention_heads
         SCREAMING_SNAKE_CASE					:			List[str]    	=       hidden_act
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       hidden_dropout_prob
         SCREAMING_SNAKE_CASE					:			Tuple    	=       attention_probs_dropout_prob
         SCREAMING_SNAKE_CASE					:			Tuple    	=       classifier_dropout_prob
         SCREAMING_SNAKE_CASE					:			Any    	=       initializer_range
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       drop_path_rate
         SCREAMING_SNAKE_CASE					:			int    	=       layer_norm_eps
         SCREAMING_SNAKE_CASE					:			Any    	=       decoder_hidden_size
         SCREAMING_SNAKE_CASE					:			List[str]    	=       kwargs.get('reshape_last_stage'    ,							lowercase__ )
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       semantic_loss_ignore_index






class 				UpperCAmelCase (       _SCREAMING_SNAKE_CASE ):




    '''simple docstring'''







    snake_case__							:   Tuple							=						version.parse("1.11" )






    @property
    def       _UpperCamelCase    (       self )							-> Mapping[str, Mapping[int, str]]:
         return OrderedDict(
             [
                 ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
             ] )






    @property
    def       _UpperCamelCase    (       self )							-> float:
         return 1E-4






    @property
    def       _UpperCamelCase    (       self )							-> int:
         return 12



