



"""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        __SCREAMING_SNAKE_CASE	:


					'''simple docstring'''





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

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



					def 		snake_case       (       self							:  Union[str, Any]     )->					Any:
								lowerCamelCase__		:		List[Any]			=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

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

								lowerCamelCase__		:		List[Any]			=self.get_config()

								return config, pixel_values, labels



					def 		snake_case       (       self							:  Any     )->					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=lowerCamelCase,  initializer_range=self.initializer_range,  mask_ratio=self.mask_ratio,  )



					def 		snake_case       (       self							:  int,  lowerCamelCase							:  int,  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  str     )->					str:
								lowerCamelCase__		:		int			=TFViTMAEModel(config=lowerCamelCase     )
								lowerCamelCase__		:		Dict			=model(lowerCamelCase,  training=lowerCamelCase     )
								self.parent.assertEqual(result.last_hidden_state.shape,  (self.batch_size, self.seq_length, self.hidden_size)     )



					def 		snake_case       (       self							:  Optional[int],  lowerCamelCase							:  Dict,  lowerCamelCase							:  Any,  lowerCamelCase							:  int     )->					int:
								lowerCamelCase__		:		int			=TFViTMAEForPreTraining(lowerCamelCase     )
								lowerCamelCase__		:		Dict			=model(lowerCamelCase,  training=lowerCamelCase     )
								# expected sequence length = num_patches
								lowerCamelCase__		:		Optional[int]			=(self.image_size // self.patch_size) ** 2
								lowerCamelCase__		:		Optional[int]			=self.patch_size**2 * self.num_channels
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, num_patches, expected_num_channels)     )

								# test greyscale images
								lowerCamelCase__		:		Tuple			=1
								lowerCamelCase__		:		Any			=TFViTMAEForPreTraining(lowerCamelCase     )

								lowerCamelCase__		:		Tuple			=floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
								lowerCamelCase__		:		List[str]			=model(lowerCamelCase,  training=lowerCamelCase     )
								lowerCamelCase__		:		Dict			=self.patch_size**2
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, num_patches, expected_num_channels)     )




					def 		snake_case       (       self							:  Optional[Any]     )->					Tuple:
								lowerCamelCase__		:		Dict			=self.prepare_config_and_inputs()
								((lowerCamelCase__)	,   (lowerCamelCase__)	,   (lowerCamelCase__))		:		Dict			=config_and_inputs
								lowerCamelCase__		:		str			={'''pixel_values''': pixel_values}
								return config, inputs_dict





@require_tf
class        __SCREAMING_SNAKE_CASE	(      lowerCAmelCase_		,   lowerCAmelCase_		,   unittest.TestCase   ):


					'''simple docstring'''





					_a   			=							(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
					_a   			=							{'feature-extraction': TFViTMAEModel} if is_tf_available() else {}

					_a   			=							False
					_a   			=							False
					_a   			=							False
					_a   			=							False



					def 		snake_case       (       self							:  int     )->					Tuple:
								lowerCamelCase__		:		Dict			=TFViTMAEModelTester(self     )
								lowerCamelCase__		:		int			=ConfigTester(self,  config_class=lowerCamelCase,  has_text_modality=lowerCamelCase,  hidden_size=37     )



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



					@unittest.skip(reason='''ViTMAE does not use inputs_embeds'''     )
					def 		snake_case       (       self							:  Optional[Any]     )->					List[str]:
								pass



					def 		snake_case       (       self							:  Any     )->					Tuple:
								lowerCamelCase__	,   lowerCamelCase__		:		Union[str, Any]			=self.model_tester.prepare_config_and_inputs_for_common()

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



					def 		snake_case       (       self							:  Optional[int]     )->					int:
								lowerCamelCase__	,   lowerCamelCase__		:		List[str]			=self.model_tester.prepare_config_and_inputs_for_common()

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

											lowerCamelCase__		:		str			=['''pixel_values''']
											self.assertListEqual(arg_names[:1],  lowerCamelCase     )



					def 		snake_case       (       self							:  Any     )->					Optional[Any]:
								lowerCamelCase__		:		int			=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase     )



					def 		snake_case       (       self							:  Union[str, Any]     )->					List[Any]:
								lowerCamelCase__		:		Any			=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*lowerCamelCase     )



					def 		snake_case       (       self							:  Union[str, Any]     )->					str:
								# make the mask reproducible
								np.random.seed(2     )

								lowerCamelCase__	,   lowerCamelCase__		:		Union[str, Any]			=self.model_tester.prepare_config_and_inputs_for_common()

								lowerCamelCase__		:		Optional[int]			=int((config.image_size // config.patch_size) ** 2     )
								lowerCamelCase__		:		Optional[Any]			=np.random.uniform(size=(self.model_tester.batch_size, num_patches)     )

								for model_class in self.all_model_classes:
											lowerCamelCase__		:		List[str]			=model_class(lowerCamelCase     )
											lowerCamelCase__		:		Tuple			=self._prepare_for_class(lowerCamelCase,  lowerCamelCase     )

											lowerCamelCase__		:		Union[str, Any]			=model(lowerCamelCase,  noise=lowerCamelCase     )

											lowerCamelCase__		:		List[Any]			=copy.deepcopy(self._prepare_for_class(lowerCamelCase,  lowerCamelCase     )     )
											lowerCamelCase__		:		Tuple			=model(**lowerCamelCase,  noise=lowerCamelCase     )
											lowerCamelCase__		:		int			=outputs_dict[0].numpy()
											lowerCamelCase__		:		Dict			=outputs_keywords[0].numpy()

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



					def 		snake_case       (       self							:  Any     )->					List[Any]:
								# make the mask reproducible
								np.random.seed(2     )

								lowerCamelCase__	,   lowerCamelCase__		:		int			=self.model_tester.prepare_config_and_inputs_for_common()

								lowerCamelCase__		:		Tuple			=int((config.image_size // config.patch_size) ** 2     )
								lowerCamelCase__		:		Optional[int]			=np.random.uniform(size=(self.model_tester.batch_size, num_patches)     )

								def prepare_numpy_arrays(lowerCamelCase							:  int     ):
											lowerCamelCase__		:		Any			={}
											for k, v in inputs_dict.items():
														if tf.is_tensor(lowerCamelCase     ):
																	lowerCamelCase__		:		Optional[Any]			=v.numpy()
														else:
																	lowerCamelCase__		:		List[str]			=np.array(lowerCamelCase     )

											return inputs_np_dict

								for model_class in self.all_model_classes:
											lowerCamelCase__		:		Tuple			=model_class(lowerCamelCase     )

											lowerCamelCase__		:		Optional[Any]			=self._prepare_for_class(lowerCamelCase,  lowerCamelCase     )
											lowerCamelCase__		:		List[Any]			=prepare_numpy_arrays(lowerCamelCase     )

											lowerCamelCase__		:		List[str]			=model(lowerCamelCase,  noise=lowerCamelCase     )
											lowerCamelCase__		:		Dict			=model(**lowerCamelCase,  noise=lowerCamelCase     )
											self.assert_outputs_same(lowerCamelCase,  lowerCamelCase     )



					def 		snake_case       (       self							:  Any,  lowerCamelCase							:  int,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  List[Any]     )->					Any:
								# make masks reproducible
								np.random.seed(2     )

								lowerCamelCase__		:		List[Any]			=int((tf_model.config.image_size // tf_model.config.patch_size) ** 2     )
								lowerCamelCase__		:		str			=np.random.uniform(size=(self.model_tester.batch_size, num_patches)     )
								lowerCamelCase__		:		Optional[int]			=tf.constant(lowerCamelCase     )

								# Add `noise` argument.
								# PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
								lowerCamelCase__		:		Optional[Any]			=tf_noise

								super().check_pt_tf_models(lowerCamelCase,  lowerCamelCase,  lowerCamelCase     )



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

								lowerCamelCase__	,   lowerCamelCase__		:		List[Any]			=self.model_tester.prepare_config_and_inputs_for_common()

								lowerCamelCase__		:		Union[str, Any]			={
								    module_member
								    for model_class in self.all_model_classes
								    for module in (import_module(model_class.__module__     ),)
								    for module_member_name in dir(lowerCamelCase     )
								    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(lowerCamelCase,  lowerCamelCase     ),)
								    if isinstance(lowerCamelCase,  lowerCamelCase     )
								    and tf.keras.layers.Layer in module_member.__bases__
								    and getattr(lowerCamelCase,  '''_keras_serializable''',  lowerCamelCase     )
								}

								lowerCamelCase__		:		Union[str, Any]			=int((config.image_size // config.patch_size) ** 2     )
								lowerCamelCase__		:		Union[str, Any]			=np.random.uniform(size=(self.model_tester.batch_size, num_patches)     )
								lowerCamelCase__		:		Optional[int]			=tf.convert_to_tensor(lowerCamelCase     )
								inputs_dict.update({'''noise''': noise}     )

								for main_layer_class in tf_main_layer_classes:
											lowerCamelCase__		:		Union[str, Any]			=main_layer_class(lowerCamelCase     )

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

											lowerCamelCase__		:		List[str]			=tf.keras.Model(lowerCamelCase,  outputs=main_layer(lowerCamelCase     )     )
											lowerCamelCase__		:		Tuple			=model(lowerCamelCase     )

											with tempfile.TemporaryDirectory() as tmpdirname:
														lowerCamelCase__		:		Dict			=os.path.join(lowerCamelCase,  '''keras_model.h5'''     )
														model.save(lowerCamelCase     )
														lowerCamelCase__		:		int			=tf.keras.models.load_model(
														    lowerCamelCase,  custom_objects={main_layer_class.__name__: main_layer_class}     )
														assert isinstance(lowerCamelCase,  tf.keras.Model     )
														lowerCamelCase__		:		List[Any]			=model(lowerCamelCase     )
														self.assert_outputs_same(lowerCamelCase,  lowerCamelCase     )



					@slow
					def 		snake_case       (       self							:  List[Any]     )->					int:
								# make mask reproducible
								np.random.seed(2     )

								lowerCamelCase__	,   lowerCamelCase__		:		Union[str, Any]			=self.model_tester.prepare_config_and_inputs_for_common()

								lowerCamelCase__		:		Union[str, Any]			=int((config.image_size // config.patch_size) ** 2     )
								lowerCamelCase__		:		str			=np.random.uniform(size=(self.model_tester.batch_size, num_patches)     )

								for model_class in self.all_model_classes:
											lowerCamelCase__		:		str			=model_class(lowerCamelCase     )
											lowerCamelCase__		:		int			=self._prepare_for_class(lowerCamelCase,  lowerCamelCase     )
											lowerCamelCase__		:		Optional[Any]			=model(lowerCamelCase,  noise=lowerCamelCase     )

											if model_class.__name__ == "TFViTMAEModel":
														lowerCamelCase__		:		Optional[int]			=outputs.last_hidden_state.numpy()
														lowerCamelCase__		:		Union[str, Any]			=0
											else:
														lowerCamelCase__		:		List[str]			=outputs.logits.numpy()
														lowerCamelCase__		:		List[Any]			=0

											with tempfile.TemporaryDirectory() as tmpdirname:
														model.save_pretrained(lowerCamelCase,  saved_model=lowerCamelCase     )
														lowerCamelCase__		:		Union[str, Any]			=model_class.from_pretrained(lowerCamelCase     )
														lowerCamelCase__		:		Any			=model(lowerCamelCase,  noise=lowerCamelCase     )

														if model_class.__name__ == "TFViTMAEModel":
																	lowerCamelCase__		:		Optional[Any]			=after_outputs['''last_hidden_state'''].numpy()
																	lowerCamelCase__		:		int			=0
														else:
																	lowerCamelCase__		:		Optional[int]			=after_outputs['''logits'''].numpy()
																	lowerCamelCase__		:		List[Any]			=0

														lowerCamelCase__		:		List[str]			=np.amax(np.abs(out_a - out_a     )     )
														self.assertLessEqual(lowerCamelCase,  1E-5     )



					def 		snake_case       (       self							:  Union[str, Any]     )->					List[Any]:
								# make mask reproducible
								np.random.seed(2     )

								lowerCamelCase__	,   lowerCamelCase__		:		Any			=self.model_tester.prepare_config_and_inputs_for_common()

								lowerCamelCase__		:		Dict			=int((config.image_size // config.patch_size) ** 2     )
								lowerCamelCase__		:		List[str]			=np.random.uniform(size=(self.model_tester.batch_size, num_patches)     )

								for model_class in self.all_model_classes:
											lowerCamelCase__		:		Dict			=model_class(lowerCamelCase     )
											lowerCamelCase__		:		Optional[int]			=self._prepare_for_class(lowerCamelCase,  lowerCamelCase     )

											lowerCamelCase__		:		Optional[Any]			=model(lowerCamelCase,  noise=lowerCamelCase     )
											lowerCamelCase__		:		Tuple			=model.get_config()
											# make sure that returned config is jsonifiable, which is required by keras
											json.dumps(lowerCamelCase     )
											lowerCamelCase__		:		Any			=model_class.from_config(model.get_config()     )
											# make sure it also accepts a normal config
											lowerCamelCase__		:		Tuple			=model_class.from_config(model.config     )
											lowerCamelCase__		:		str			=new_model(lowerCamelCase     )  # Build model
											new_model.set_weights(model.get_weights()     )
											lowerCamelCase__		:		Optional[Any]			=new_model(lowerCamelCase,  noise=lowerCamelCase     )

											self.assert_outputs_same(lowerCamelCase,  lowerCamelCase     )



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



					@unittest.skip(reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load'''     )
					def 		snake_case       (       self							:  Tuple     )->					Dict:
								pass




					@slow
					def 		snake_case       (       self							:  Optional[int]     )->					List[Any]:
								lowerCamelCase__		:		Union[str, Any]			=TFViTMAEModel.from_pretrained('''google/vit-base-patch16-224'''     )
								self.assertIsNotNone(lowerCamelCase     )






def      snake_case__   (		):
			"""simple docstring"""
			lowerCamelCase__		:		str			=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''	)
			return image

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


					'''simple docstring'''





					@cached_property
					def 		snake_case       (       self							:  int     )->					Tuple:
								return ViTImageProcessor.from_pretrained('''facebook/vit-mae-base'''     ) if is_vision_available() else None




					@slow
					def 		snake_case       (       self							:  str     )->					Optional[Any]:
								# make random mask reproducible across the PT and TF model
								np.random.seed(2     )

								lowerCamelCase__		:		Optional[Any]			=TFViTMAEForPreTraining.from_pretrained('''facebook/vit-mae-base'''     )

								lowerCamelCase__		:		str			=self.default_image_processor
								lowerCamelCase__		:		Optional[int]			=prepare_img()
								lowerCamelCase__		:		Dict			=image_processor(images=lowerCamelCase,  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)
								lowerCamelCase__		:		Optional[int]			=ViTMAEConfig()
								lowerCamelCase__		:		Dict			=int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2     )
								lowerCamelCase__		:		List[str]			=np.random.uniform(size=(1, num_patches)     )

								# forward pass
								lowerCamelCase__		:		List[str]			=model(**lowerCamelCase,  noise=lowerCamelCase     )

								# verify the logits
								lowerCamelCase__		:		Any			=tf.convert_to_tensor([1, 196, 768]     )
								self.assertEqual(outputs.logits.shape,  lowerCamelCase     )

								lowerCamelCase__		:		Tuple			=tf.convert_to_tensor(
								    [[-0.0_548, -1.7_023, -0.9_325], [0.3_721, -0.5_670, -0.2_233], [0.8_235, -1.3_878, -0.3_524]]     )

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









"""simple docstring"""






from functools import lru_cache






def      snake_case__   (		__lowerCamelCase				:    int	):
			"""simple docstring"""
			lowerCamelCase__		:		Optional[Any]			=2
			lowerCamelCase__		:		Optional[int]			=set()
			while i * i <= n:
						if n % i:
									i += 1
						else:
									n //= i
									factors.add(__lowerCamelCase	)
			if n > 1:
						factors.add(__lowerCamelCase	)
			return factors






@lru_cache
def      snake_case__   (		__lowerCamelCase				:    int	):
			"""simple docstring"""
			return len(unique_prime_factors(__lowerCamelCase	)	)






def      snake_case__   (		__lowerCamelCase				:    list	):
			"""simple docstring"""
			return len(set(__lowerCamelCase	)	) in (0, 1)






def      snake_case__   (		__lowerCamelCase				:    int	):
			"""simple docstring"""
			lowerCamelCase__		:		Tuple			=2

			while True:
						# Increment each value of a generated range
						lowerCamelCase__		:		Tuple			=[base + i for i in range(__lowerCamelCase	)]

						# Run elements through out unique_prime_factors function
						# Append our target number to the end.
						lowerCamelCase__		:		Optional[Any]			=[upf_len(__lowerCamelCase	) for x in group]
						checker.append(__lowerCamelCase	)

						# If all numbers in the list are equal, return the group variable.
						if equality(__lowerCamelCase	):
									return group

						# Increment our base variable by 1
						base += 1






def      snake_case__   (		__lowerCamelCase				:    int = 4	):
			"""simple docstring"""
			lowerCamelCase__		:		List[Any]			=run(__lowerCamelCase	)
			return results[0] if len(__lowerCamelCase	) else None


if __name__ == "__main__":
		print(solution())




