





'''simple docstring'''






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




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

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




      def 							lowercase__       (     self     :					Any    )							-> Any:
            _lowerCAmelCase					=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

            _lowerCAmelCase					=				None
            if self.use_labels:
                  _lowerCAmelCase					=				ids_tensor([self.batch_size]     ,     self.type_sequence_label_size    )

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

            return config, pixel_values, labels




      def 							lowercase__       (     self     :					Optional[int]     ,     __snake_case     :					str     ,     __snake_case     :					List[str]     ,     __snake_case     :					Dict    )							-> List[Any]:
            _lowerCAmelCase					=				FlaxBeitModel(config=__snake_case    )
            _lowerCAmelCase					=				model(__snake_case    )
            self.parent.assertEqual(result.last_hidden_state.shape     ,     (self.batch_size, self.seq_length, self.hidden_size)    )




      def 							lowercase__       (     self     :					Optional[Any]     ,     __snake_case     :					Optional[Any]     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					List[Any]    )							-> Union[str, Any]:
            _lowerCAmelCase					=				FlaxBeitForMaskedImageModeling(config=__snake_case    )
            _lowerCAmelCase					=				model(__snake_case    )
            self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.seq_length - 1, self.vocab_size)    )




      def 							lowercase__       (     self     :					List[str]     ,     __snake_case     :					Tuple     ,     __snake_case     :					str     ,     __snake_case     :					str    )							-> Dict:
            _lowerCAmelCase					=				self.type_sequence_label_size
            _lowerCAmelCase					=				FlaxBeitForImageClassification(config=__snake_case    )
            _lowerCAmelCase					=				model(__snake_case    )
            self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.type_sequence_label_size)    )

            # test greyscale images
            _lowerCAmelCase					=				1
            _lowerCAmelCase					=				FlaxBeitForImageClassification(__snake_case    )

            _lowerCAmelCase					=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
            _lowerCAmelCase					=				model(__snake_case    )




      def 							lowercase__       (     self     :					str    )							-> Dict:
            _lowerCAmelCase					=				self.prepare_config_and_inputs()
            (
                (
                _lowerCAmelCase
            )	,				(
                _lowerCAmelCase
            )	,				(
                _lowerCAmelCase
            )	,				
            )					=				config_and_inputs
            _lowerCAmelCase					=				{"""pixel_values""": pixel_values}
            return config, inputs_dict




@require_flax
class 					UpperCAmelCase	(     snake_case_				,						unittest.TestCase  ):
      _lowercase:     List[str]			=  (
          (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
      )




      def 							lowercase__       (     self     :					Optional[int]    )							-> None:
            _lowerCAmelCase					=				FlaxBeitModelTester(self    )
            _lowerCAmelCase					=				ConfigTester(self     ,     config_class=__snake_case     ,     has_text_modality=__snake_case     ,     hidden_size=37    )




      def 							lowercase__       (     self     :					Optional[Any]    )							-> Union[str, Any]:
            self.config_tester.run_common_tests()




      def 							lowercase__       (     self     :					Tuple    )							-> List[Any]:
            _lowerCAmelCase	,				_lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs_for_common()

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

                  _lowerCAmelCase					=				["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]     ,     __snake_case    )




      def 							lowercase__       (     self     :					Optional[Any]    )							-> Tuple:
            _lowerCAmelCase	,				_lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                  with self.subTest(model_class.__name__    ):
                        _lowerCAmelCase					=				self._prepare_for_class(__snake_case     ,     __snake_case    )
                        _lowerCAmelCase					=				model_class(__snake_case    )

                        @jax.jit
                        def model_jitted(__snake_case     :					List[str]     ,     **__snake_case     :					List[Any]    ):
                              return model(pixel_values=__snake_case     ,     **__snake_case    )

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

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




      def 							lowercase__       (     self     :					List[str]    )							-> List[Any]:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*__snake_case    )




      def 							lowercase__       (     self     :					Optional[Any]    )							-> int:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_masked_lm(*__snake_case    )




      def 							lowercase__       (     self     :					int    )							-> Optional[int]:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*__snake_case    )




      @slow
      def 							lowercase__       (     self     :					Dict    )							-> Union[str, Any]:
            for model_class_name in self.all_model_classes:
                  _lowerCAmelCase					=				model_class_name.from_pretrained("""microsoft/beit-base-patch16-224"""    )
                  _lowerCAmelCase					=				model(np.ones((1, 3, 2_24, 2_24)    )    )
                  self.assertIsNotNone(__snake_case    )




def   UpperCamelCase__ (						):





      """simple docstring"""


      _lowerCAmelCase					=				Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
      return image





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




      @cached_property
      def 							lowercase__       (     self     :					int    )							-> List[str]:
            return BeitImageProcessor.from_pretrained("""microsoft/beit-base-patch16-224"""    ) if is_vision_available() else None




      @slow
      def 							lowercase__       (     self     :					List[str]    )							-> Tuple:
            _lowerCAmelCase					=				FlaxBeitForMaskedImageModeling.from_pretrained("""microsoft/beit-base-patch16-224-pt22k"""    )

            _lowerCAmelCase					=				self.default_image_processor
            _lowerCAmelCase					=				prepare_img()
            _lowerCAmelCase					=				image_processor(images=__snake_case     ,     return_tensors="""np"""    ).pixel_values

            # prepare bool_masked_pos
            _lowerCAmelCase					=				np.ones((1, 1_96)     ,     dtype=__snake_case    )

            # forward pass
            _lowerCAmelCase					=				model(pixel_values=__snake_case     ,     bool_masked_pos=__snake_case    )
            _lowerCAmelCase					=				outputs.logits

            # verify the logits
            _lowerCAmelCase					=				(1, 1_96, 81_92)
            self.assertEqual(logits.shape     ,     __snake_case    )

            _lowerCAmelCase					=				np.array(
                [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]    )

            self.assertTrue(np.allclose(logits[bool_masked_pos][:3, :3]     ,     __snake_case     ,     atol=1E-2    )    )




      @slow
      def 							lowercase__       (     self     :					Optional[Any]    )							-> List[Any]:
            _lowerCAmelCase					=				FlaxBeitForImageClassification.from_pretrained("""microsoft/beit-base-patch16-224"""    )

            _lowerCAmelCase					=				self.default_image_processor
            _lowerCAmelCase					=				prepare_img()
            _lowerCAmelCase					=				image_processor(images=__snake_case     ,     return_tensors="""np"""    )

            # forward pass
            _lowerCAmelCase					=				model(**__snake_case    )
            _lowerCAmelCase					=				outputs.logits

            # verify the logits
            _lowerCAmelCase					=				(1, 10_00)
            self.assertEqual(logits.shape     ,     __snake_case    )

            _lowerCAmelCase					=				np.array([-1.23_85, -1.09_87, -1.01_08]    )

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

            _lowerCAmelCase					=				2_81
            self.assertEqual(logits.argmax(-1    ).item()     ,     __snake_case    )




      @slow
      def 							lowercase__       (     self     :					List[str]    )							-> Any:
            _lowerCAmelCase					=				FlaxBeitForImageClassification.from_pretrained("""microsoft/beit-large-patch16-224-pt22k-ft22k"""    )

            _lowerCAmelCase					=				self.default_image_processor
            _lowerCAmelCase					=				prepare_img()
            _lowerCAmelCase					=				image_processor(images=__snake_case     ,     return_tensors="""np"""    )

            # forward pass
            _lowerCAmelCase					=				model(**__snake_case    )
            _lowerCAmelCase					=				outputs.logits

            # verify the logits
            _lowerCAmelCase					=				(1, 2_18_41)
            self.assertEqual(logits.shape     ,     __snake_case    )

            _lowerCAmelCase					=				np.array([1.68_81, -0.27_87, 0.59_01]    )

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

            _lowerCAmelCase					=				23_96
            self.assertEqual(logits.argmax(-1    ).item()     ,     __snake_case    )













'''simple docstring'''






def   UpperCamelCase__ (						lowerCAmelCase    ):





      """simple docstring"""


      if not isinstance(lowerCAmelCase				,      lowerCAmelCase    ):
            raise ValueError("""Input must be an integer"""    )
      if input_num <= 0:
            raise ValueError("""Input must be positive"""    )
      return sum(
          divisor for divisor in range(1				,      input_num // 2 + 1    ) if input_num % divisor == 0    )


if __name__ == "__main__":
      import doctest

      doctest.testmod()






