






import inspect
import unittest

import numpy as np

from transformers import BeitConfig
from transformers.testing_utils import require_flax, require_vision, slow
from transformers.utils import cached_property, is_flax_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor


if is_flax_available():
 import jax

 from transformers import FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling, FlaxBeitModel

if is_vision_available():
 from PIL import Image

 from transformers import BeitImageProcessor







class 	_UpperCAmelCase  (  unittest.TestCase		):




      """simple docstring"""



      def __init__(   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase=100    ,	_lowerCAmelCase=13    ,	_lowerCAmelCase=30    ,	_lowerCAmelCase=2    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=True    ,	_lowerCAmelCase=True    ,	_lowerCAmelCase=32    ,	_lowerCAmelCase=5    ,	_lowerCAmelCase=4    ,	_lowerCAmelCase=37    ,	_lowerCAmelCase="gelu"    ,	_lowerCAmelCase=0.1    ,	_lowerCAmelCase=0.1    ,	_lowerCAmelCase=10    ,	_lowerCAmelCase=0.02    ,	_lowerCAmelCase=3    ,	):



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

          # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
          lowerCAmelCase__   :Optional[Any]      							=				(image_size // patch_size) ** 2
          lowerCAmelCase__   :List[str]      							=				num_patches + 1
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Any      							=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

          lowerCAmelCase__   :Optional[int]      							=				None
          if self.use_labels:
              lowerCAmelCase__   :Optional[Any]      							=				ids_tensor([self.batch_size]    ,	self.type_sequence_label_size		)

          lowerCAmelCase__   :Tuple      							=				BeitConfig(
              vocab_size=self.vocab_size    ,	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, labels
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :List[str]      							=				FlaxBeitModel(config=_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				model(_lowerCAmelCase		)
          self.parent.assertEqual(result.last_hidden_state.shape    ,	(self.batch_size, self.seq_length, self.hidden_size)		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				FlaxBeitForMaskedImageModeling(config=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[int]      							=				model(_lowerCAmelCase		)
          self.parent.assertEqual(result.logits.shape    ,	(self.batch_size, self.seq_length - 1, self.vocab_size)		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :List[Any]      							=				self.type_sequence_label_size
          lowerCAmelCase__   :List[Any]      							=				FlaxBeitForImageClassification(config=_lowerCAmelCase		)
          lowerCAmelCase__   :Tuple      							=				model(_lowerCAmelCase		)
          self.parent.assertEqual(result.logits.shape    ,	(self.batch_size, self.type_sequence_label_size)		)

          # test greyscale images
          lowerCAmelCase__   :List[str]      							=				1
          lowerCAmelCase__   :Optional[Any]      							=				FlaxBeitForImageClassification(_lowerCAmelCase		)

          lowerCAmelCase__   :int      							=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
          lowerCAmelCase__   :Any      							=				model(_lowerCAmelCase		)



      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Dict      							=				self.prepare_config_and_inputs()
          (
              (
              lowerCAmelCase__
          )							,(
              lowerCAmelCase__
          )							,(
              lowerCAmelCase__
          )							,
          )   :int      							=				config_and_inputs
          lowerCAmelCase__   :Optional[int]      							=				{"pixel_values": pixel_values}
          return config, inputs_dict







@require_flax
class 	_UpperCAmelCase  (  _A				,							unittest.TestCase		):




      """simple docstring"""



      A          =    (
          (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
      )
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :List[Any]      							=				FlaxBeitModelTester(self		)
          lowerCAmelCase__   :Optional[Any]      							=				ConfigTester(self    ,	config_class=_lowerCAmelCase    ,	has_text_modality=_lowerCAmelCase    ,	hidden_size=37		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          self.config_tester.run_common_tests()
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__							,lowerCAmelCase__   :Tuple      							=				self.model_tester.prepare_config_and_inputs_for_common()

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

              lowerCAmelCase__   :Union[str, Any]      							=				["pixel_values"]
              self.assertListEqual(arg_names[:1]    ,	_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__							,lowerCAmelCase__   :Dict      							=				self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__		):
                  lowerCAmelCase__   :Optional[int]      							=				self._prepare_for_class(_lowerCAmelCase    ,	_lowerCAmelCase		)
                  lowerCAmelCase__   :Optional[int]      							=				model_class(_lowerCAmelCase		)

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

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

                  with self.subTest("JIT Disabled"		):
                      with jax.disable_jit():
                          lowerCAmelCase__   :str      							=				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		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :List[Any]      							=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Any      							=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*_lowerCAmelCase		)



      @slow
      def 			snake_case_      (   self		):



          '''simple docstring'''
          for model_class_name in self.all_model_classes:
              lowerCAmelCase__   :Optional[Any]      							=				model_class_name.from_pretrained("microsoft/beit-base-patch16-224"		)
              lowerCAmelCase__   :Union[str, Any]      							=				model(np.ones((1, 3, 224, 224)		)		)
              self.assertIsNotNone(_lowerCAmelCase		)





def 				snake_case__							(							):
    lowerCAmelCase__   :int      							=				Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
    return image







@require_vision
@require_flax
class 	_UpperCAmelCase  (  unittest.TestCase		):




      """simple docstring"""



      @cached_property
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"		) if is_vision_available() else None
      @slow
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"		)

          lowerCAmelCase__   :int      							=				self.default_image_processor
          lowerCAmelCase__   :str      							=				prepare_img()
          lowerCAmelCase__   :Tuple      							=				image_processor(images=_lowerCAmelCase    ,	return_tensors="np"		).pixel_values

          # prepare bool_masked_pos
          lowerCAmelCase__   :int      							=				np.ones((1, 196)    ,	dtype=_lowerCAmelCase		)

          # forward pass
          lowerCAmelCase__   :Union[str, Any]      							=				model(pixel_values=_lowerCAmelCase    ,	bool_masked_pos=_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				outputs.logits

          # verify the logits
          lowerCAmelCase__   :Dict      							=				(1, 196, 8_192)
          self.assertEqual(logits.shape    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :List[str]      							=				np.array(
              [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]		)

          self.assertTrue(np.allclose(logits[bool_masked_pos][:3, :3]    ,	_lowerCAmelCase    ,	atol=1e-2		)		)
      @slow
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"		)

          lowerCAmelCase__   :List[str]      							=				self.default_image_processor
          lowerCAmelCase__   :Dict      							=				prepare_img()
          lowerCAmelCase__   :Tuple      							=				image_processor(images=_lowerCAmelCase    ,	return_tensors="np"		)

          # forward pass
          lowerCAmelCase__   :Union[str, Any]      							=				model(**_lowerCAmelCase		)
          lowerCAmelCase__   :Any      							=				outputs.logits

          # verify the logits
          lowerCAmelCase__   :Union[str, Any]      							=				(1, 1_000)
          self.assertEqual(logits.shape    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :Dict      							=				np.array([-1.2385, -1.0987, -1.0108]		)

          self.assertTrue(np.allclose(logits[0, :3]    ,	_lowerCAmelCase    ,	atol=1e-4		)		)

          lowerCAmelCase__   :List[Any]      							=				281
          self.assertEqual(logits.argmax(-1		).item()    ,	_lowerCAmelCase		)



      @slow
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"		)

          lowerCAmelCase__   :Optional[int]      							=				self.default_image_processor
          lowerCAmelCase__   :Tuple      							=				prepare_img()
          lowerCAmelCase__   :List[Any]      							=				image_processor(images=_lowerCAmelCase    ,	return_tensors="np"		)

          # forward pass
          lowerCAmelCase__   :Optional[Any]      							=				model(**_lowerCAmelCase		)
          lowerCAmelCase__   :str      							=				outputs.logits

          # verify the logits
          lowerCAmelCase__   :List[Any]      							=				(1, 21_841)
          self.assertEqual(logits.shape    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :Any      							=				np.array([1.6881, -0.2787, 0.5901]		)

          self.assertTrue(np.allclose(logits[0, :3]    ,	_lowerCAmelCase    ,	atol=1e-4		)		)

          lowerCAmelCase__   :Optional[int]      							=				2_396
          self.assertEqual(logits.argmax(-1		).item()    ,	_lowerCAmelCase		)











from ...configuration_utils import PretrainedConfig
from ...utils import logging


_a    :   Dict        =			logging.get_logger(__name__)

_a    :   Optional[Any]        =			{
    """caidas/swin2sr-classicalsr-x2-64""": (
        """https://huggingface.co/caidas/swin2sr-classicalsr-x2-64/resolve/main/config.json"""
    ),
}







class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      A          =    '''swin2sr'''

      A          =    {
          '''hidden_size''': '''embed_dim''',
          '''num_attention_heads''': '''num_heads''',
          '''num_hidden_layers''': '''num_layers''',
      }
      def __init__(   self    ,	_lowerCAmelCase=64    ,	_lowerCAmelCase=1    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=180    ,	_lowerCAmelCase=[6, 6, 6, 6, 6, 6]    ,	_lowerCAmelCase=[6, 6, 6, 6, 6, 6]    ,	_lowerCAmelCase=8    ,	_lowerCAmelCase=2.0    ,	_lowerCAmelCase=True    ,	_lowerCAmelCase=0.0    ,	_lowerCAmelCase=0.0    ,	_lowerCAmelCase=0.1    ,	_lowerCAmelCase="gelu"    ,	_lowerCAmelCase=False    ,	_lowerCAmelCase=0.02    ,	_lowerCAmelCase=1e-5    ,	_lowerCAmelCase=2    ,	_lowerCAmelCase=1.0    ,	_lowerCAmelCase="1conv"    ,	_lowerCAmelCase="pixelshuffle"    ,	**_lowerCAmelCase    ,	):



          '''simple docstring'''
          super().__init__(**_lowerCAmelCase		)

          lowerCAmelCase__   :Union[str, Any]      							=				image_size
          lowerCAmelCase__   :List[Any]      							=				patch_size
          lowerCAmelCase__   :Any      							=				num_channels
          lowerCAmelCase__   :Union[str, Any]      							=				embed_dim
          lowerCAmelCase__   :Optional[int]      							=				depths
          lowerCAmelCase__   :int      							=				len(_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				num_heads
          lowerCAmelCase__   :str      							=				window_size
          lowerCAmelCase__   :List[str]      							=				mlp_ratio
          lowerCAmelCase__   :List[Any]      							=				qkv_bias
          lowerCAmelCase__   :Union[str, Any]      							=				hidden_dropout_prob
          lowerCAmelCase__   :Optional[int]      							=				attention_probs_dropout_prob
          lowerCAmelCase__   :Dict      							=				drop_path_rate
          lowerCAmelCase__   :Tuple      							=				hidden_act
          lowerCAmelCase__   :Dict      							=				use_absolute_embeddings
          lowerCAmelCase__   :Tuple      							=				layer_norm_eps
          lowerCAmelCase__   :Dict      							=				initializer_range
          lowerCAmelCase__   :Optional[int]      							=				upscale
          lowerCAmelCase__   :Optional[Any]      							=				img_range
          lowerCAmelCase__   :List[str]      							=				resi_connection
          lowerCAmelCase__   :Union[str, Any]      							=				upsampler



