


'''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 						_snake_case   :




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

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




						def 			snake_case__  ( self):
									UpperCAmelCase__					:   Any       				=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

									UpperCAmelCase__					:   Any       				=       None
									if self.use_labels:
												UpperCAmelCase__					:   List[str]       				=       ids_tensor([self.batch_size]							,							self.type_sequence_label_size)

									UpperCAmelCase__					:   Any       				=       self.get_config()

									return config, pixel_values, labels




						def 			snake_case__  ( self):
									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=_lowercase							,							initializer_range=self.initializer_range							,							mask_ratio=self.mask_ratio							,							)




						def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
									UpperCAmelCase__					:   int       				=       TFViTMAEModel(config=_lowercase)
									UpperCAmelCase__					:   List[Any]       				=       model(_lowercase							,							training=_lowercase)
									self.parent.assertEqual(result.last_hidden_state.shape							,							(self.batch_size, self.seq_length, self.hidden_size))




						def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
									UpperCAmelCase__					:   List[str]       				=       TFViTMAEForPreTraining(_lowercase)
									UpperCAmelCase__					:   List[Any]       				=       model(_lowercase							,							training=_lowercase)
									# expected sequence length = num_patches
									UpperCAmelCase__					:   List[str]       				=       (self.image_size // self.patch_size) ** 2
									UpperCAmelCase__					:   Optional[Any]       				=       self.patch_size**2 * self.num_channels
									self.parent.assertEqual(result.logits.shape							,							(self.batch_size, num_patches, expected_num_channels))

									# test greyscale images
									UpperCAmelCase__					:   List[str]       				=       1
									UpperCAmelCase__					:   Dict       				=       TFViTMAEForPreTraining(_lowercase)

									UpperCAmelCase__					:   Dict       				=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
									UpperCAmelCase__					:   Optional[Any]       				=       model(_lowercase							,							training=_lowercase)
									UpperCAmelCase__					:   Dict       				=       self.patch_size**2
									self.parent.assertEqual(result.logits.shape							,							(self.batch_size, num_patches, expected_num_channels))




						def 			snake_case__  ( self):
									UpperCAmelCase__					:   Optional[int]       				=       self.prepare_config_and_inputs()
									((UpperCAmelCase__)				,    (UpperCAmelCase__)				,    (UpperCAmelCase__))					:   Dict       				=       config_and_inputs
									UpperCAmelCase__					:   int       				=       {"""pixel_values""": pixel_values}
									return config, inputs_dict


@require_tf
class 						_snake_case   (       _lowerCAmelCase    ,				_lowerCAmelCase    ,				unittest.TestCase		):
						lowerCAmelCase :int						=   (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
						lowerCAmelCase :Optional[Any]						=   {"feature-extraction": TFViTMAEModel} if is_tf_available() else {}

						lowerCAmelCase :Any						=   False
						lowerCAmelCase :Dict						=   False
						lowerCAmelCase :Dict						=   False
						lowerCAmelCase :int						=   False




						def 			snake_case__  ( self):
									UpperCAmelCase__					:   Optional[Any]       				=       TFViTMAEModelTester(self)
									UpperCAmelCase__					:   Union[str, Any]       				=       ConfigTester(self							,							config_class=_lowercase							,							has_text_modality=_lowercase							,							hidden_size=37)




						def 			snake_case__  ( self):
									self.config_tester.run_common_tests()




						@unittest.skip(reason="""ViTMAE does not use inputs_embeds""")
						def 			snake_case__  ( self):
									pass




						def 			snake_case__  ( self):
									UpperCAmelCase__				,    UpperCAmelCase__					:   int       				=       self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
												UpperCAmelCase__					:   Union[str, Any]       				=       model_class(_lowercase)
												self.assertIsInstance(model.get_input_embeddings()							,							(tf.keras.layers.Layer))
												UpperCAmelCase__					:   Optional[int]       				=       model.get_output_embeddings()
												self.assertTrue(x is None or isinstance(_lowercase							,							tf.keras.layers.Layer))




						def 			snake_case__  ( self):
									UpperCAmelCase__				,    UpperCAmelCase__					:   Any       				=       self.model_tester.prepare_config_and_inputs_for_common()

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

												UpperCAmelCase__					:   Optional[Any]       				=       ["""pixel_values"""]
												self.assertListEqual(arg_names[:1]							,							_lowercase)




						def 			snake_case__  ( self):
									UpperCAmelCase__					:   List[str]       				=       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_lowercase)




						def 			snake_case__  ( self):
									UpperCAmelCase__					:   Optional[Any]       				=       self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*_lowercase)




						def 			snake_case__  ( self):
									# make the mask reproducible
									np.random.seed(2)

									UpperCAmelCase__				,    UpperCAmelCase__					:   Optional[Any]       				=       self.model_tester.prepare_config_and_inputs_for_common()

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

									for model_class in self.all_model_classes:
												UpperCAmelCase__					:   Optional[int]       				=       model_class(_lowercase)
												UpperCAmelCase__					:   int       				=       self._prepare_for_class(_lowercase							,							_lowercase)

												UpperCAmelCase__					:   int       				=       model(_lowercase							,							noise=_lowercase)

												UpperCAmelCase__					:   str       				=       copy.deepcopy(self._prepare_for_class(_lowercase							,							_lowercase))
												UpperCAmelCase__					:   Union[str, Any]       				=       model(**_lowercase							,							noise=_lowercase)
												UpperCAmelCase__					:   Optional[Any]       				=       outputs_dict[0].numpy()
												UpperCAmelCase__					:   Optional[Any]       				=       outputs_keywords[0].numpy()

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




						def 			snake_case__  ( self):
									# make the mask reproducible
									np.random.seed(2)

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

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

									def prepare_numpy_arrays(_lowerCamelCase):
												UpperCAmelCase__					:   str       				=       {}
												for k, v in inputs_dict.items():
															if tf.is_tensor(_lowercase):
																		UpperCAmelCase__					:   Dict       				=       v.numpy()
															else:
																		UpperCAmelCase__					:   Optional[int]       				=       np.array(_lowercase)

												return inputs_np_dict

									for model_class in self.all_model_classes:
												UpperCAmelCase__					:   Dict       				=       model_class(_lowercase)

												UpperCAmelCase__					:   List[str]       				=       self._prepare_for_class(_lowercase							,							_lowercase)
												UpperCAmelCase__					:   int       				=       prepare_numpy_arrays(_lowercase)

												UpperCAmelCase__					:   List[str]       				=       model(_lowercase							,							noise=_lowercase)
												UpperCAmelCase__					:   str       				=       model(**_lowercase							,							noise=_lowercase)
												self.assert_outputs_same(_lowercase							,							_lowercase)




						def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
									# make masks reproducible
									np.random.seed(2)

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

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

									super().check_pt_tf_models(_lowercase							,							_lowercase							,							_lowercase)




						def 			snake_case__  ( self):
									# make mask reproducible
									np.random.seed(2)

									UpperCAmelCase__				,    UpperCAmelCase__					:   Optional[int]       				=       self.model_tester.prepare_config_and_inputs_for_common()

									UpperCAmelCase__					:   int       				=       {
									    module_member
									    for model_class in self.all_model_classes
									    for module in (import_module(model_class.__module__),)
									    for module_member_name in dir(_lowercase)
									    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(_lowercase							,							_lowercase),)
									    if isinstance(_lowercase							,							_lowercase)
									    and tf.keras.layers.Layer in module_member.__bases__
									    and getattr(_lowercase							,							"""_keras_serializable"""							,							_lowercase)
									}

									UpperCAmelCase__					:   int       				=       int((config.image_size // config.patch_size) ** 2)
									UpperCAmelCase__					:   Tuple       				=       np.random.uniform(size=(self.model_tester.batch_size, num_patches))
									UpperCAmelCase__					:   List[Any]       				=       tf.convert_to_tensor(_lowercase)
									inputs_dict.update({"""noise""": noise})

									for main_layer_class in tf_main_layer_classes:
												UpperCAmelCase__					:   List[Any]       				=       main_layer_class(_lowercase)

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

												UpperCAmelCase__					:   Any       				=       tf.keras.Model(_lowercase							,							outputs=main_layer(_lowercase))
												UpperCAmelCase__					:   Any       				=       model(_lowercase)

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




						@slow
						def 			snake_case__  ( self):
									# make mask reproducible
									np.random.seed(2)

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

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

									for model_class in self.all_model_classes:
												UpperCAmelCase__					:   Tuple       				=       model_class(_lowercase)
												UpperCAmelCase__					:   str       				=       self._prepare_for_class(_lowercase							,							_lowercase)
												UpperCAmelCase__					:   Dict       				=       model(_lowercase							,							noise=_lowercase)

												if model_class.__name__ == "TFViTMAEModel":
															UpperCAmelCase__					:   Dict       				=       outputs.last_hidden_state.numpy()
															UpperCAmelCase__					:   int       				=       0
												else:
															UpperCAmelCase__					:   Any       				=       outputs.logits.numpy()
															UpperCAmelCase__					:   List[str]       				=       0

												with tempfile.TemporaryDirectory() as tmpdirname:
															model.save_pretrained(_lowercase							,							saved_model=_lowercase)
															UpperCAmelCase__					:   Any       				=       model_class.from_pretrained(_lowercase)
															UpperCAmelCase__					:   List[Any]       				=       model(_lowercase							,							noise=_lowercase)

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

															UpperCAmelCase__					:   Optional[int]       				=       np.amax(np.abs(out_a - out_a))
															self.assertLessEqual(_lowercase							,							1e-5)




						def 			snake_case__  ( self):
									# make mask reproducible
									np.random.seed(2)

									UpperCAmelCase__				,    UpperCAmelCase__					:   List[str]       				=       self.model_tester.prepare_config_and_inputs_for_common()

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

									for model_class in self.all_model_classes:
												UpperCAmelCase__					:   Dict       				=       model_class(_lowercase)
												UpperCAmelCase__					:   Any       				=       self._prepare_for_class(_lowercase							,							_lowercase)

												UpperCAmelCase__					:   Tuple       				=       model(_lowercase							,							noise=_lowercase)
												UpperCAmelCase__					:   str       				=       model.get_config()
												# make sure that returned config is jsonifiable, which is required by keras
												json.dumps(_lowercase)
												UpperCAmelCase__					:   int       				=       model_class.from_config(model.get_config())
												# make sure it also accepts a normal config
												UpperCAmelCase__					:   Dict       				=       model_class.from_config(model.config)
												UpperCAmelCase__					:   Dict       				=       new_model(_lowercase)  # Build model
												new_model.set_weights(model.get_weights())
												UpperCAmelCase__					:   int       				=       new_model(_lowercase							,							noise=_lowercase)

												self.assert_outputs_same(_lowercase							,							_lowercase)




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




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




						@slow
						def 			snake_case__  ( self):
									UpperCAmelCase__					:   Optional[int]       				=       TFViTMAEModel.from_pretrained("""google/vit-base-patch16-224""")
									self.assertIsNotNone(_lowercase)







def       _UpperCamelCase (						):
			UpperCAmelCase__					:   Optional[int]       				=       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
			return image


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




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




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

									UpperCAmelCase__					:   str       				=       TFViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base""")

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

									# forward pass
									UpperCAmelCase__					:   Dict       				=       model(**_lowercase							,							noise=_lowercase)

									# verify the logits
									UpperCAmelCase__					:   Union[str, Any]       				=       tf.convert_to_tensor([1, 196, 768])
									self.assertEqual(outputs.logits.shape							,							_lowercase)

									UpperCAmelCase__					:   str       				=       tf.convert_to_tensor(
									    [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]])

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



'''simple docstring'''



import colorsys

from PIL import Image  # type: ignore







def       _UpperCamelCase (						UpperCamelCase__	, UpperCamelCase__	, UpperCamelCase__   ):
			UpperCAmelCase__					:   Optional[Any]       				=       x
			UpperCAmelCase__					:   Optional[int]       				=       y
			for step in range(UpperCamelCase__   ):  # noqa: B007
						UpperCAmelCase__					:   List[str]       				=       a * a - b * b + x
						UpperCAmelCase__					:   Optional[int]       				=       2 * a * b + y
						UpperCAmelCase__					:   int       				=       a_new

						# divergence happens for all complex number with an absolute value
						# greater than 4
						if a * a + b * b > 4:
									break
			return step / (max_step - 1)







def       _UpperCamelCase (						UpperCamelCase__   ):
			if distance == 1:
						return (0, 0, 0)
			else:
						return (2_5_5, 2_5_5, 2_5_5)







def       _UpperCamelCase (						UpperCamelCase__   ):
			if distance == 1:
						return (0, 0, 0)
			else:
						return tuple(round(i * 2_5_5   ) for i in colorsys.hsv_to_rgb(UpperCamelCase__	, 1	, 1   )   )







def       _UpperCamelCase (						UpperCamelCase__ = 8_0_0	, UpperCamelCase__ = 6_0_0	, UpperCamelCase__ = -0.6	, UpperCamelCase__ = 0	, UpperCamelCase__ = 3.2	, UpperCamelCase__ = 5_0	, UpperCamelCase__ = True	, ):
			UpperCAmelCase__					:   str       				=       Image.new("""RGB"""	, (image_width, image_height)   )
			UpperCAmelCase__					:   Optional[int]       				=       img.load()

			# loop through the image-coordinates
			for image_x in range(UpperCamelCase__   ):
						for image_y in range(UpperCamelCase__   ):
									# determine the figure-coordinates based on the image-coordinates
									UpperCAmelCase__					:   Union[str, Any]       				=       figure_width / image_width * image_height
									UpperCAmelCase__					:   Union[str, Any]       				=       figure_center_x + (image_x / image_width - 0.5) * figure_width
									UpperCAmelCase__					:   Tuple       				=       figure_center_y + (image_y / image_height - 0.5) * figure_height

									UpperCAmelCase__					:   List[str]       				=       get_distance(UpperCamelCase__	, UpperCamelCase__	, UpperCamelCase__   )

									# color the corresponding pixel based on the selected coloring-function
									if use_distance_color_coding:
												UpperCAmelCase__					:   Any       				=       get_color_coded_rgb(UpperCamelCase__   )
									else:
												UpperCAmelCase__					:   Tuple       				=       get_black_and_white_rgb(UpperCamelCase__   )

			return img


if __name__ == "__main__":
		import doctest

		doctest.testmod()

		# colored version, full figure
		__A		=get_image()

		# uncomment for colored version, different section, zoomed in
		# img = get_image(figure_center_x = -0.6, figure_center_y = -0.4,
		# figure_width = 0.8)

		# uncomment for black and white version, full figure
		# img = get_image(use_distance_color_coding = False)

		# uncomment to save the image
		# img.save("mandelbrot.png")

		img.show()