




'''simple docstring'''





from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
   from PIL import Image

   from transformers import ViTImageProcessor

class   _a							:






     def __init__(					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE=13   ,_SCREAMING_SNAKE_CASE=30   ,_SCREAMING_SNAKE_CASE=2   ,_SCREAMING_SNAKE_CASE=3   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=32   ,_SCREAMING_SNAKE_CASE=2   ,_SCREAMING_SNAKE_CASE=4   ,_SCREAMING_SNAKE_CASE=37   ,_SCREAMING_SNAKE_CASE="gelu"   ,_SCREAMING_SNAKE_CASE=0.1   ,_SCREAMING_SNAKE_CASE=0.1   ,_SCREAMING_SNAKE_CASE=10   ,_SCREAMING_SNAKE_CASE=0.0_2   ,_SCREAMING_SNAKE_CASE=3   ,_SCREAMING_SNAKE_CASE=0.6   ,_SCREAMING_SNAKE_CASE=None   ,)				->     Union[str, Any]:
          _snake_case				= parent
          _snake_case				= batch_size
          _snake_case				= image_size
          _snake_case				= patch_size
          _snake_case				= num_channels
          _snake_case				= is_training
          _snake_case				= use_labels
          _snake_case				= hidden_size
          _snake_case				= num_hidden_layers
          _snake_case				= num_attention_heads
          _snake_case				= intermediate_size
          _snake_case				= hidden_act
          _snake_case				= hidden_dropout_prob
          _snake_case				= attention_probs_dropout_prob
          _snake_case				= type_sequence_label_size
          _snake_case				= initializer_range
          _snake_case				= mask_ratio
          _snake_case				= scope

          # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
          # (we add 1 for the [CLS] token)
          _snake_case				= (image_size // patch_size) ** 2
          _snake_case				= int(math.ceil((1 - mask_ratio) * (num_patches + 1)       )       )






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

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

          _snake_case				= self.get_config()

          return config, pixel_values, labels






     def 							_lowercase      (					self       )				->     List[str]:
          return ViTMAEConfig(
              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   ,decoder_hidden_size=self.hidden_size   ,decoder_num_hidden_layers=self.num_hidden_layers   ,decoder_num_attention_heads=self.num_attention_heads   ,decoder_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   ,mask_ratio=self.mask_ratio   ,)






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     List[Any]:
          _snake_case				= TFViTMAEModel(config=__a       )
          _snake_case				= model(__a   ,training=__a       )
          self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size)       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Optional[int]:
          _snake_case				= TFViTMAEForPreTraining(__a       )
          _snake_case				= model(__a   ,training=__a       )
          # expected sequence length = num_patches
          _snake_case				= (self.image_size // self.patch_size) ** 2
          _snake_case				= self.patch_size**2 * self.num_channels
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, num_patches, expected_num_channels)       )

          # test greyscale images
          _snake_case				= 1
          _snake_case				= TFViTMAEForPreTraining(__a       )

          _snake_case				= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
          _snake_case				= model(__a   ,training=__a       )
          _snake_case				= self.patch_size**2
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, num_patches, expected_num_channels)       )






     def 							_lowercase      (					self       )				->     Tuple:
          _snake_case				= self.prepare_config_and_inputs()
          (_snake_case)				= config_and_inputs
          _snake_case				= {"""pixel_values""": pixel_values}
          return config, inputs_dict







@require_tf
class   _a							(		a__   ,   a__   ,   unittest.TestCase				):
     SCREAMING_SNAKE_CASE_				:       Dict       							=		(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
     SCREAMING_SNAKE_CASE_				:       List[str]       							=		{"""feature-extraction""": TFViTMAEModel} if is_tf_available() else {}

     SCREAMING_SNAKE_CASE_				:       str       							=		False
     SCREAMING_SNAKE_CASE_				:       Optional[Any]       							=		False
     SCREAMING_SNAKE_CASE_				:       Optional[Any]       							=		False
     SCREAMING_SNAKE_CASE_				:       Tuple       							=		False






     def 							_lowercase      (					self       )				->     Union[str, Any]:
          _snake_case				= TFViTMAEModelTester(self       )
          _snake_case				= ConfigTester(self   ,config_class=__a   ,has_text_modality=__a   ,hidden_size=37       )






     def 							_lowercase      (					self       )				->     str:
          self.config_tester.run_common_tests()






     @unittest.skip(reason="ViTMAE does not use inputs_embeds"       )
     def 							_lowercase      (					self       )				->     Optional[int]:
          pass






     def 							_lowercase      (					self       )				->     Any:
          _snake_case				= self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
               _snake_case				= model_class(__a       )
               self.assertIsInstance(model.get_input_embeddings()   ,(tf.keras.layers.Layer)       )
               _snake_case				= model.get_output_embeddings()
               self.assertTrue(x is None or isinstance(__a   ,tf.keras.layers.Layer       )       )






     def 							_lowercase      (					self       )				->     Tuple:
          _snake_case				= self.model_tester.prepare_config_and_inputs_for_common()

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

               _snake_case				= ["""pixel_values"""]
               self.assertListEqual(arg_names[:1]   ,__a       )






     def 							_lowercase      (					self       )				->     List[str]:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__a       )






     def 							_lowercase      (					self       )				->     List[str]:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_pretraining(*__a       )






     def 							_lowercase      (					self       )				->     int:
          # make the mask reproducible
          np.random.seed(2       )

          _snake_case				= self.model_tester.prepare_config_and_inputs_for_common()

          _snake_case				= int((config.image_size // config.patch_size) ** 2       )
          _snake_case				= np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

          for model_class in self.all_model_classes:
               _snake_case				= model_class(__a       )
               _snake_case				= self._prepare_for_class(__a   ,__a       )

               _snake_case				= model(__a   ,noise=__a       )

               _snake_case				= copy.deepcopy(self._prepare_for_class(__a   ,__a       )       )
               _snake_case				= model(**__a   ,noise=__a       )
               _snake_case				= outputs_dict[0].numpy()
               _snake_case				= outputs_keywords[0].numpy()

               self.assertLess(np.sum(np.abs(output_dict - output_keywords       )       )   ,1e-6       )






     def 							_lowercase      (					self       )				->     Optional[int]:
          # make the mask reproducible
          np.random.seed(2       )

          _snake_case				= self.model_tester.prepare_config_and_inputs_for_common()

          _snake_case				= int((config.image_size // config.patch_size) ** 2       )
          _snake_case				= np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

          def prepare_numpy_arrays(_SCREAMING_SNAKE_CASE       ):
               _snake_case				= {}
               for k, v in inputs_dict.items():
                    if tf.is_tensor(__a       ):
                         _snake_case				= v.numpy()
                    else:
                         _snake_case				= np.array(__a       )

               return inputs_np_dict

          for model_class in self.all_model_classes:
               _snake_case				= model_class(__a       )

               _snake_case				= self._prepare_for_class(__a   ,__a       )
               _snake_case				= prepare_numpy_arrays(__a       )

               _snake_case				= model(__a   ,noise=__a       )
               _snake_case				= model(**__a   ,noise=__a       )
               self.assert_outputs_same(__a   ,__a       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Union[str, Any]:
          # make masks reproducible
          np.random.seed(2       )

          _snake_case				= int((tf_model.config.image_size // tf_model.config.patch_size) ** 2       )
          _snake_case				= np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )
          _snake_case				= tf.constant(__a       )

          # Add `noise` argument.
          # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
          _snake_case				= tf_noise

          super().check_pt_tf_models(__a   ,__a   ,__a       )






     def 							_lowercase      (					self       )				->     Optional[Any]:
          # make mask reproducible
          np.random.seed(2       )

          _snake_case				= self.model_tester.prepare_config_and_inputs_for_common()

          _snake_case				= {
              module_member
              for model_class in self.all_model_classes
              for module in (import_module(model_class.__module__       ),)
              for module_member_name in dir(__a       )
              if module_member_name.endswith("MainLayer"       )
              # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
              and module_member_name[: -len("MainLayer"       )] == model_class.__name__[: -len("Model"       )]
              for module_member in (getattr(__a   ,__a       ),)
              if isinstance(__a   ,__a       )
              and tf.keras.layers.Layer in module_member.__bases__
              and getattr(__a   ,"_keras_serializable"   ,__a       )
          }

          _snake_case				= int((config.image_size // config.patch_size) ** 2       )
          _snake_case				= np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )
          _snake_case				= tf.convert_to_tensor(__a       )
          inputs_dict.update({"noise": noise}       )

          for main_layer_class in tf_main_layer_classes:
               _snake_case				= main_layer_class(__a       )

               _snake_case				= {
                   name: tf.keras.Input(tensor.shape[1:]   ,dtype=tensor.dtype       ) for name, tensor in inputs_dict.items()
               }

               _snake_case				= tf.keras.Model(__a   ,outputs=main_layer(__a       )       )
               _snake_case				= model(__a       )

               with tempfile.TemporaryDirectory() as tmpdirname:
                    _snake_case				= os.path.join(__a   ,"keras_model.h5"       )
                    model.save(__a       )
                    _snake_case				= tf.keras.models.load_model(
                        __a   ,custom_objects={main_layer_class.__name__: main_layer_class}       )
                    assert isinstance(__a   ,tf.keras.Model       )
                    _snake_case				= model(__a       )
                    self.assert_outputs_same(__a   ,__a       )






     @slow
     def 							_lowercase      (					self       )				->     Dict:
          # make mask reproducible
          np.random.seed(2       )

          _snake_case				= self.model_tester.prepare_config_and_inputs_for_common()

          _snake_case				= int((config.image_size // config.patch_size) ** 2       )
          _snake_case				= np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

          for model_class in self.all_model_classes:
               _snake_case				= model_class(__a       )
               _snake_case				= self._prepare_for_class(__a   ,__a       )
               _snake_case				= model(__a   ,noise=__a       )

               if model_class.__name__ == "TFViTMAEModel":
                    _snake_case				= outputs.last_hidden_state.numpy()
                    _snake_case				= 0
               else:
                    _snake_case				= outputs.logits.numpy()
                    _snake_case				= 0

               with tempfile.TemporaryDirectory() as tmpdirname:
                    model.save_pretrained(__a   ,saved_model=__a       )
                    _snake_case				= model_class.from_pretrained(__a       )
                    _snake_case				= model(__a   ,noise=__a       )

                    if model_class.__name__ == "TFViTMAEModel":
                         _snake_case				= after_outputs["""last_hidden_state"""].numpy()
                         _snake_case				= 0
                    else:
                         _snake_case				= after_outputs["""logits"""].numpy()
                         _snake_case				= 0

                    _snake_case				= np.amax(np.abs(out_a - out_a       )       )
                    self.assertLessEqual(__a   ,1e-5       )






     def 							_lowercase      (					self       )				->     str:
          # make mask reproducible
          np.random.seed(2       )

          _snake_case				= self.model_tester.prepare_config_and_inputs_for_common()

          _snake_case				= int((config.image_size // config.patch_size) ** 2       )
          _snake_case				= np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

          for model_class in self.all_model_classes:
               _snake_case				= model_class(__a       )
               _snake_case				= self._prepare_for_class(__a   ,__a       )

               _snake_case				= model(__a   ,noise=__a       )
               _snake_case				= model.get_config()
               # make sure that returned config is jsonifiable, which is required by keras
               json.dumps(__a       )
               _snake_case				= model_class.from_config(model.get_config()       )
               # make sure it also accepts a normal config
               _snake_case				= model_class.from_config(model.config       )
               _snake_case				= new_model(__a       )  # Build model
               new_model.set_weights(model.get_weights()       )
               _snake_case				= new_model(__a   ,noise=__a       )

               self.assert_outputs_same(__a   ,__a       )






     @unittest.skip(
         reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."       )
     def 							_lowercase      (					self       )				->     Dict:
          pass






     @unittest.skip(reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"       )
     def 							_lowercase      (					self       )				->     int:
          pass






     @slow
     def 							_lowercase      (					self       )				->     List[str]:
          _snake_case				= TFViTMAEModel.from_pretrained("google/vit-base-patch16-224"       )
          self.assertIsNotNone(__a       )







def       __a      (  )		->					Dict:

     """simple docstring"""




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





@require_tf
@require_vision
class   _a							(		unittest.TestCase				):






     @cached_property
     def 							_lowercase      (					self       )				->     Dict:
          return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"       ) if is_vision_available() else None






     @slow
     def 							_lowercase      (					self       )				->     Dict:
          # make random mask reproducible across the PT and TF model
          np.random.seed(2       )

          _snake_case				= TFViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"       )

          _snake_case				= self.default_image_processor
          _snake_case				= prepare_img()
          _snake_case				= image_processor(images=__a   ,return_tensors="tf"       )

          # prepare a noise vector that will be also used for testing the TF model
          # (this way we can ensure that the PT and TF models operate on the same inputs)
          _snake_case				= ViTMAEConfig()
          _snake_case				= int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2       )
          _snake_case				= np.random.uniform(size=(1, num_patches)       )

          # forward pass
          _snake_case				= model(**__a   ,noise=__a       )

          # verify the logits
          _snake_case				= tf.convert_to_tensor([1, 196, 768]       )
          self.assertEqual(outputs.logits.shape   ,__a       )

          _snake_case				= tf.convert_to_tensor(
              [[-0.0_5_4_8, -1.7_0_2_3, -0.9_3_2_5], [0.3_7_2_1, -0.5_6_7_0, -0.2_2_3_3], [0.8_2_3_5, -1.3_8_7_8, -0.3_5_2_4]]       )

          tf.debugging.assert_near(outputs.logits[0, :3, :3]   ,__a   ,atol=1e-4       )








'''simple docstring'''







def        lowerCamelCase__				(      __lowercase				):
     if not isinstance(__lowercase     ,   __lowercase				):
          snake_case			:      int        =			F'''Input value of [number={number}] must be an integer'''
          raise TypeError(__lowercase				)
     if number < 0:
          return False
     snake_case			:      Optional[int]        =			number * number
     while number > 0:
          if number % 10 != number_square % 10:
               return False
          number //= 10
          number_square //= 10
     return True


if __name__ == "__main__":
      import doctest

      doctest.testmod()



