






"""simple docstring"""



import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch
						from torch import nn

						from transformers import ViTMAEForPreTraining, ViTMAEModel
						from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import ViTImageProcessor







class   lowercase__     :


			def __init__( self	:    str  ,			snake_case__	:    Dict  ,			snake_case__	:    Tuple=13  ,			snake_case__	:    Optional[int]=30  ,			snake_case__	:    Optional[int]=2  ,			snake_case__	:    str=3  ,			snake_case__	:    Optional[int]=True  ,			snake_case__	:    List[Any]=True  ,			snake_case__	:    Optional[int]=32  ,			snake_case__	:    Tuple=5  ,			snake_case__	:    Tuple=4  ,			snake_case__	:    Union[str, Any]=37  ,			snake_case__	:    Union[str, Any]="gelu"  ,			snake_case__	:    Optional[int]=0.1  ,			snake_case__	:    Optional[int]=0.1  ,			snake_case__	:    Union[str, Any]=10  ,			snake_case__	:    Tuple=0.02  ,			snake_case__	:    Optional[int]=3  ,			snake_case__	:    str=0.6  ,			snake_case__	:    List[str]=None  ,			):
										lowerCamelCase_			:				int										=parent
										lowerCamelCase_			:				Any										=batch_size
										lowerCamelCase_			:				Any										=image_size
										lowerCamelCase_			:				Optional[int]										=patch_size
										lowerCamelCase_			:				int										=num_channels
										lowerCamelCase_			:				Optional[int]										=is_training
										lowerCamelCase_			:				str										=use_labels
										lowerCamelCase_			:				List[Any]										=hidden_size
										lowerCamelCase_			:				str										=num_hidden_layers
										lowerCamelCase_			:				List[Any]										=num_attention_heads
										lowerCamelCase_			:				str										=intermediate_size
										lowerCamelCase_			:				Any										=hidden_act
										lowerCamelCase_			:				int										=hidden_dropout_prob
										lowerCamelCase_			:				Tuple										=attention_probs_dropout_prob
										lowerCamelCase_			:				Union[str, Any]										=type_sequence_label_size
										lowerCamelCase_			:				str										=initializer_range
										lowerCamelCase_			:				Union[str, Any]										=mask_ratio
										lowerCamelCase_			:				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_			:				str										=int(math.ceil((1 - mask_ratio) * (num_patches + 1)   )   )


			def       UpperCAmelCase__		( self	:    List[Any]   ):
										lowerCamelCase_			:				Optional[int]										=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

										lowerCamelCase_			:				int										=None
										if self.use_labels:
																	lowerCamelCase_			:				Tuple										=ids_tensor([self.batch_size]  ,			self.type_sequence_label_size   )

										lowerCamelCase_			:				Any										=self.get_config()

										return config, pixel_values, labels


			def       UpperCAmelCase__		( self	:    Union[str, Any]   ):
										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  ,			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  ,			mask_ratio=self.mask_ratio  ,			)


			def       UpperCAmelCase__		( self	:    int  ,			snake_case__	:    int  ,			snake_case__	:    int  ,			snake_case__	:    Tuple   ):
										lowerCamelCase_			:				int										=ViTMAEModel(config=snake_case__   )
										model.to(snake_case__   )
										model.eval()
										lowerCamelCase_			:				Dict										=model(snake_case__   )
										self.parent.assertEqual(result.last_hidden_state.shape  ,			(self.batch_size, self.seq_length, self.hidden_size)   )


			def       UpperCAmelCase__		( self	:    Union[str, Any]  ,			snake_case__	:    List[str]  ,			snake_case__	:    Union[str, Any]  ,			snake_case__	:    Dict   ):
										lowerCamelCase_			:				Optional[Any]										=ViTMAEForPreTraining(snake_case__   )
										model.to(snake_case__   )
										model.eval()
										lowerCamelCase_			:				Any										=model(snake_case__   )
										lowerCamelCase_			:				List[Any]										=(self.image_size // self.patch_size) ** 2
										lowerCamelCase_			:				Dict										=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_			:				int										=1
										lowerCamelCase_			:				str										=ViTMAEForPreTraining(snake_case__   )
										model.to(snake_case__   )
										model.eval()
										lowerCamelCase_			:				Dict										=floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
										lowerCamelCase_			:				List[str]										=model(snake_case__   )
										lowerCamelCase_			:				Dict										=self.patch_size**2
										self.parent.assertEqual(result.logits.shape  ,			(self.batch_size, num_patches, expected_num_channels)   )


			def       UpperCAmelCase__		( self	:    str   ):
										lowerCamelCase_			:				Tuple										=self.prepare_config_and_inputs()
										lowerCamelCase_			:				int										=config_and_inputs
										lowerCamelCase_			:				List[str]										={"pixel_values": pixel_values}
										return config, inputs_dict







@require_torch
class   lowercase__     (				snake_case__,  snake_case__,  unittest.TestCase  ):
			_UpperCAmelCase						:Any       =					(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
			_UpperCAmelCase						:int       =					{"""feature-extraction""": ViTMAEModel} if is_torch_available() else {}

			_UpperCAmelCase						:int       =					False
			_UpperCAmelCase						:List[str]       =					False
			_UpperCAmelCase						:Dict       =					False
			_UpperCAmelCase						:Dict       =					False


			def       UpperCAmelCase__		( self	:    Union[str, Any]   ):
										lowerCamelCase_			:				int										=ViTMAEModelTester(self   )
										lowerCamelCase_			:				Any										=ConfigTester(self  ,			config_class=snake_case__  ,			has_text_modality=snake_case__  ,			hidden_size=37   )


			def       UpperCAmelCase__		( self	:    str   ):
										self.config_tester.run_common_tests()


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


			def       UpperCAmelCase__		( self	:    Dict   ):
										lowerCamelCase_			:				Dict										=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																	lowerCamelCase_			:				Optional[Any]										=model_class(snake_case__   )
																	self.assertIsInstance(model.get_input_embeddings()  ,			(nn.Module)   )
																	lowerCamelCase_			:				Optional[Any]										=model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(snake_case__  ,			nn.Linear   )   )


			def       UpperCAmelCase__		( self	:    Dict   ):
										lowerCamelCase_			:				Any										=self.model_tester.prepare_config_and_inputs_for_common()

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

																	lowerCamelCase_			:				Tuple										=["pixel_values"]
																	self.assertListEqual(arg_names[:1]  ,			snake_case__   )


			def       UpperCAmelCase__		( self	:    Dict   ):
										lowerCamelCase_			:				str										=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*snake_case__   )


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


			def       UpperCAmelCase__		( self	:    Union[str, Any]  ,			snake_case__	:    Tuple  ,			snake_case__	:    List[Any]  ,			snake_case__	:    int   ):
										# make masks reproducible
										np.random.seed(2   )

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

										# Add `noise` argument.
										# PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
										lowerCamelCase_			:				Union[str, Any]										=pt_noise

										super().check_pt_tf_models(snake_case__  ,			snake_case__  ,			snake_case__   )


			def       UpperCAmelCase__		( self	:    Optional[int]   ):
										lowerCamelCase_			:				Optional[Any]										=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																	lowerCamelCase_			:				Optional[int]										=model_class(snake_case__   )
																	model.to(snake_case__   )
																	model.eval()
																	# make random mask reproducible
																	torch.manual_seed(2   )
																	with torch.no_grad():
																								lowerCamelCase_			:				List[str]										=model(**self._prepare_for_class(snake_case__  ,			snake_case__   )   )

																	lowerCamelCase_			:				Optional[int]										=outputs[0].cpu().numpy()
																	lowerCamelCase_			:				List[str]										=0

																	with tempfile.TemporaryDirectory() as tmpdirname:
																								model.save_pretrained(snake_case__   )
																								lowerCamelCase_			:				Union[str, Any]										=model_class.from_pretrained(snake_case__   )
																								model.to(snake_case__   )
																								# make random mask reproducible
																								torch.manual_seed(2   )
																								with torch.no_grad():
																															lowerCamelCase_			:				Union[str, Any]										=model(**self._prepare_for_class(snake_case__  ,			snake_case__   )   )

																								# Make sure we don't have nans
																								lowerCamelCase_			:				str										=after_outputs[0].cpu().numpy()
																								lowerCamelCase_			:				Tuple										=0
																								lowerCamelCase_			:				Dict										=np.amax(np.abs(out_a - out_a   )   )
																								self.assertLessEqual(snake_case__  ,			1E-5   )


			@unittest.skip(
			    reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."   )
			def       UpperCAmelCase__		( self	:    List[Any]   ):
										pass


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


			@unittest.skip(
			    reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."   )
			def       UpperCAmelCase__		( self	:    List[str]   ):
										pass


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


			@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."   )
			def       UpperCAmelCase__		( self	:    int   ):
										pass


			@slow
			def       UpperCAmelCase__		( self	:    Optional[Any]   ):
										for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowerCamelCase_			:				Tuple										=ViTMAEModel.from_pretrained(snake_case__   )
																	self.assertIsNotNone(snake_case__   )


def 					_snake_case							(	)							->      str:
							lowerCamelCase_			:				List[str]										=Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
							return image







@require_torch
@require_vision
class   lowercase__     (				unittest.TestCase  ):


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


			@slow
			def       UpperCAmelCase__		( self	:    List[str]   ):
										# make random mask reproducible across the PT and TF model
										np.random.seed(2   )

										lowerCamelCase_			:				Tuple										=ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"   ).to(snake_case__   )

										lowerCamelCase_			:				Optional[int]										=self.default_image_processor
										lowerCamelCase_			:				Union[str, Any]										=prepare_img()
										lowerCamelCase_			:				Dict										=image_processor(images=snake_case__  ,			return_tensors="pt"   ).to(snake_case__   )

										# 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_			:				Union[str, Any]										=ViTMAEConfig()
										lowerCamelCase_			:				int										=int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2   )
										lowerCamelCase_			:				int										=np.random.uniform(size=(1, num_patches)   )

										# forward pass
										with torch.no_grad():
																	lowerCamelCase_			:				Tuple										=model(**snake_case__  ,			noise=torch.from_numpy(snake_case__   ).to(device=snake_case__   )   )

										# verify the logits
										lowerCamelCase_			:				Tuple										=torch.Size((1, 196, 768)   )
										self.assertEqual(outputs.logits.shape  ,			snake_case__   )

										lowerCamelCase_			:				Union[str, Any]										=torch.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]]   )

										self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]  ,			expected_slice.to(snake_case__   )  ,			atol=1E-4   )   )







import unittest

from transformers.utils.backbone_utils import (
    BackboneMixin,
    get_aligned_output_features_output_indices,
    verify_out_features_out_indices,
)



class 						_UpperCAmelCase      (     unittest.TestCase							):






		def _snake_case		(						self				:     Union[str, Any]):
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						["a", "b", "c"]

				# Defaults to last layer if both are None
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Dict				     =						get_aligned_output_features_output_indices(UpperCAmelCase   ,  UpperCAmelCase   ,  UpperCAmelCase)
				self.assertEqual(UpperCAmelCase   ,  ["c"])
				self.assertEqual(UpperCAmelCase   ,  [2])

				# Out indices set to match out features
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Tuple				     =						get_aligned_output_features_output_indices(["a", "c"]   ,  UpperCAmelCase   ,  UpperCAmelCase)
				self.assertEqual(UpperCAmelCase   ,  ["a", "c"])
				self.assertEqual(UpperCAmelCase   ,  [0, 2])

				# Out features set to match out indices
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:int				     =						get_aligned_output_features_output_indices(UpperCAmelCase   ,  [0, 2]   ,  UpperCAmelCase)
				self.assertEqual(UpperCAmelCase   ,  ["a", "c"])
				self.assertEqual(UpperCAmelCase   ,  [0, 2])

				# Out features selected from negative indices
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Any				     =						get_aligned_output_features_output_indices(UpperCAmelCase   ,  [-3, -1]   ,  UpperCAmelCase)
				self.assertEqual(UpperCAmelCase   ,  ["a", "c"])
				self.assertEqual(UpperCAmelCase   ,  [-3, -1])






		def _snake_case		(						self				:     str):
				# Stage names must be set
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["a", "b"]   ,  (0, 1)   ,  UpperCAmelCase)

				# Out features must be a list
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(("a", "b")   ,  (0, 1)   ,  ["a", "b"])

				# Out features must be a subset of stage names
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["a", "b"]   ,  (0, 1)   ,  ["a"])

				# Out indices must be a list or tuple
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(UpperCAmelCase   ,  0   ,  ["a", "b"])

				# Out indices must be a subset of stage names
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(UpperCAmelCase   ,  (0, 1)   ,  ["a"])

				# Out features and out indices must be the same length
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["a", "b"]   ,  (0,)   ,  ["a", "b", "c"])

				# Out features should match out indices
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["a", "b"]   ,  (0, 2)   ,  ["a", "b", "c"])

				# Out features and out indices should be in order
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["b", "a"]   ,  (0, 1)   ,  ["a", "b"])

				# Check passes with valid inputs
				verify_out_features_out_indices(["a", "b", "d"]   ,  (0, 1, -1)   ,  ["a", "b", "c", "d"])







		def _snake_case		(						self				:     Tuple):
				SCREAMING_SNAKE_CASE_				:int				     =						BackboneMixin()

				SCREAMING_SNAKE_CASE_				:List[str]				     =						["a", "b", "c"]
				SCREAMING_SNAKE_CASE_				:str				     =						["a", "c"]
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						[0, 2]

				# Check that the output features and indices are set correctly
				self.assertEqual(backbone.out_features   ,  ["a", "c"])
				self.assertEqual(backbone.out_indices   ,  [0, 2])

				# Check out features and indices are updated correctly
				SCREAMING_SNAKE_CASE_				:Tuple				     =						["a", "b"]
				self.assertEqual(backbone.out_features   ,  ["a", "b"])
				self.assertEqual(backbone.out_indices   ,  [0, 1])

				SCREAMING_SNAKE_CASE_				:List[str]				     =						[-3, -1]
				self.assertEqual(backbone.out_features   ,  ["a", "c"])
				self.assertEqual(backbone.out_indices   ,  [-3, -1])






