



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







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

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







    def 			lowercase_       (      self :  Union[str, Any]	):
      a						:					Any        =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

      a						:					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=__snake_case						,			initializer_range=self.initializer_range						,			)

      return config, pixel_values







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

      # test greyscale images
      a						:					Dict        =     1
      a						:					List[str]        =     FlaxViTForImageClassification(__snake_case	)

      a						:					Tuple        =     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
      a						:					Union[str, Any]        =     model(__snake_case	)







    def 			lowercase_       (      self :  Dict	):
      a						:					int        =     self.prepare_config_and_inputs()
      (
          a				
      )						:					Any        =     config_and_inputs
      a						:					Optional[int]        =     {'''pixel_values''': pixel_values}
      return config, inputs_dict

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







    def 			lowercase_       (      self :  Tuple	):
      a						:					List[str]        =     FlaxViTModelTester(self	)
      a						:					str        =     ConfigTester(self						,			config_class=__snake_case						,			has_text_modality=__snake_case						,			hidden_size=37	)







    def 			lowercase_       (      self :  List[Any]	):
      self.config_tester.run_common_tests()







    def 			lowercase_       (      self :  Union[str, Any]	):
      a						:					Optional[Any]        =     self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__snake_case	)







    def 			lowercase_       (      self :  List[Any]	):
      a						:					Optional[Any]        =     self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*__snake_case	)







    def 			lowercase_       (      self :  Any	):
      a						:					Dict        =     self.model_tester.prepare_config_and_inputs_for_common()

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

        a						:					Dict        =     ['''pixel_values''']
        self.assertListEqual(arg_names[:1]						,			__snake_case	)







    def 			lowercase_       (      self :  Any	):
      a						:					str        =     self.model_tester.prepare_config_and_inputs_for_common()

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

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

          with self.subTest('JIT Enabled'	):
            a						:					List[str]        =     model_jitted(**__snake_case	).to_tuple()

          with self.subTest('JIT Disabled'	):
            with jax.disable_jit():
              a						:					Tuple        =     model_jitted(**__snake_case	).to_tuple()

          self.assertEqual(len(__snake_case	)						,			len(__snake_case	)	)
          for jitted_output, output in zip(__snake_case						,			__snake_case	):
            self.assertEqual(jitted_output.shape						,			output.shape	)







    @slow
    def 			lowercase_       (      self :  Union[str, Any]	):
      for model_class_name in self.all_model_classes:
        a						:					Union[str, Any]        =     model_class_name.from_pretrained('google/vit-base-patch16-224'	)
        a						:					Tuple        =     model(np.ones((1, 3, 2_24, 2_24)	)	)
        self.assertIsNotNone(__snake_case	)




"""simple docstring"""
import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
_lowerCAmelCase    :str     =					object()

# For specifying empty leaf dict `{}`
_lowerCAmelCase    :str     =					object()





def  lowerCamelCase_						(UpperCamelCase__		:		List[str]     ,       UpperCamelCase__		:		int							):
    _UpperCAmelCase							:  Dict											=    tuple((re.compile(x + '''$'''							) for x in qs)							)
    for i in range(len(UpperCamelCase__							) - len(UpperCamelCase__							) + 1							):
        _UpperCAmelCase							:  str											=    [x.match(UpperCamelCase__							) for x, y in zip(UpperCamelCase__     ,       ks[i:]							)]
        if matches and all(UpperCamelCase__							):
            return True
    return False





def  lowerCamelCase_						(UpperCamelCase__		:		List[str]							):
    def replace(UpperCamelCase__		:		List[str]     ,       UpperCamelCase__		:		Tuple							):
        for rule, replacement in rules:
            if _match(UpperCamelCase__     ,       UpperCamelCase__							):
                return replacement
        return val

    return replace





def  lowerCamelCase_						():
    return [
        # embeddings
        (("transformer", "wpe", "embedding"), P('''mp'''     ,       UpperCamelCase__							)),
        (("transformer", "wte", "embedding"), P('''mp'''     ,       UpperCamelCase__							)),
        # atention
        (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(UpperCamelCase__     ,       '''mp'''							)),
        (("attention", "out_proj", "kernel"), P('''mp'''     ,       UpperCamelCase__							)),
        (("attention", "out_proj", "bias"), None),
        # mlp
        (("mlp", "c_fc", "kernel"), P(UpperCamelCase__     ,       '''mp'''							)),
        (("mlp", "c_fc", "bias"), P('''mp'''							)),
        (("mlp", "c_proj", "kernel"), P('''mp'''     ,       UpperCamelCase__							)),
        (("mlp", "c_proj", "bias"), None),
        # layer norms
        ((r"ln_\d+", "bias"), None),
        ((r"\d+", r"ln_\d+", "scale"), None),
        (("ln_f", "bias"), None),
        (("ln_f", "scale"), None),
    ]





def  lowerCamelCase_						(UpperCamelCase__		:		str							):
    _UpperCAmelCase							:  List[str]											=    _get_partition_rules()
    _UpperCAmelCase							:  List[str]											=    _replacement_rules(UpperCamelCase__							)
    _UpperCAmelCase							:  List[Any]											=    {k: _unmatched for k in flatten_dict(UpperCamelCase__							)}
    _UpperCAmelCase							:  int											=    {k: replace(UpperCamelCase__     ,       UpperCamelCase__							) for k, v in initd.items()}
    assert _unmatched not in result.values(), "Incomplete partition spec."
    return freeze(unflatten_dict(UpperCamelCase__							)							)


