





"""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 					__UpperCAmelCase							:
				def __init__(		self    :			Any	, a_    :			int	, a_    :			int=13	, a_    :			Any=30	, a_    :			Any=2	, a_    :			Tuple=3	, a_    :			Union[str, Any]=True	, a_    :			Optional[int]=True	, a_    :			List[str]=32	, a_    :			Dict=5	, a_    :			List[Any]=4	, a_    :			Optional[Any]=37	, a_    :			Dict="gelu"	, a_    :			Union[str, Any]=0.1	, a_    :			str=0.1	, a_    :			Optional[int]=10	, a_    :			Optional[int]=0.02	, a_    :			Union[str, Any]=3	, a_    :			List[Any]=0.6	, a_    :			Any=None	, )	->	Tuple:



							'''simple docstring'''



							a__  :	str		 =			parent
							a__  :	Union[str, Any]		 =			batch_size
							a__  :	Union[str, Any]		 =			image_size
							a__  :	Dict		 =			patch_size
							a__  :	Optional[Any]		 =			num_channels
							a__  :	List[Any]		 =			is_training
							a__  :	Dict		 =			use_labels
							a__  :	Optional[Any]		 =			hidden_size
							a__  :	int		 =			num_hidden_layers
							a__  :	List[Any]		 =			num_attention_heads
							a__  :	Any		 =			intermediate_size
							a__  :	Tuple		 =			hidden_act
							a__  :	List[Any]		 =			hidden_dropout_prob
							a__  :	Tuple		 =			attention_probs_dropout_prob
							a__  :	List[Any]		 =			type_sequence_label_size
							a__  :	List[Any]		 =			initializer_range
							a__  :	Union[str, Any]		 =			mask_ratio
							a__  :	Tuple		 =			scope

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



							'''simple docstring'''



							a__  :	List[str]		 =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

							a__  :	int		 =			self.get_config()

							return config, pixel_values, labels
				def       UpperCAmelCase      (		self    :			Tuple    )	->	List[str]:



							'''simple docstring'''



							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=a_	, initializer_range=self.initializer_range	, mask_ratio=self.mask_ratio	, )
				def       UpperCAmelCase      (		self    :			Optional[Any]	, a_    :			str	, a_    :			Tuple	, a_    :			Optional[int]    )	->	Any:



							'''simple docstring'''



							a__  :	Tuple		 =			ViTMAEModel(config=a_    )
							model.to(a_    )
							model.eval()
							a__  :	Union[str, Any]		 =			model(a_    )
							self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, self.seq_length, self.hidden_size)    )
				def       UpperCAmelCase      (		self    :			Union[str, Any]	, a_    :			str	, a_    :			int	, a_    :			Optional[Any]    )	->	List[str]:



							'''simple docstring'''



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






				def       UpperCAmelCase      (		self    :			Dict    )	->	Tuple:



							'''simple docstring'''



							a__  :	str		 =			self.prepare_config_and_inputs()
							a__ ,	a__ ,	a__  :	List[Any]		 =			config_and_inputs
							a__  :	Optional[Any]		 =			{"pixel_values": pixel_values}
							return config, inputs_dict



@require_torch
class 					__UpperCAmelCase							(						_UpperCamelCase						,  _UpperCamelCase						,  unittest.TestCase  ):
				__lowerCamelCase				:	Optional[Any]				     =				(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
				__lowerCamelCase				:	List[Any]				     =				{"feature-extraction": ViTMAEModel} if is_torch_available() else {}

				__lowerCamelCase				:	List[str]				     =				False
				__lowerCamelCase				:	Tuple				     =				False
				__lowerCamelCase				:	Tuple				     =				False
				__lowerCamelCase				:	int				     =				False
				def       UpperCAmelCase      (		self    :			Dict    )	->	int:



							'''simple docstring'''



							a__  :	List[str]		 =			ViTMAEModelTester(self    )
							a__  :	Any		 =			ConfigTester(self	, config_class=a_	, has_text_modality=a_	, hidden_size=37    )
				def       UpperCAmelCase      (		self    :			Union[str, Any]    )	->	int:



							'''simple docstring'''



							self.config_tester.run_common_tests()
				@unittest.skip(reason="ViTMAE does not use inputs_embeds"    )
				def       UpperCAmelCase      (		self    :			Optional[Any]    )	->	Optional[int]:



							'''simple docstring'''



							pass
				def       UpperCAmelCase      (		self    :			List[Any]    )	->	List[str]:



							'''simple docstring'''



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

							for model_class in self.all_model_classes:
										a__  :	Dict		 =			model_class(a_    )
										self.assertIsInstance(model.get_input_embeddings()	, (nn.Module)    )
										a__  :	Any		 =			model.get_output_embeddings()
										self.assertTrue(x is None or isinstance(a_	, nn.Linear    )    )
				def       UpperCAmelCase      (		self    :			Any    )	->	int:



							'''simple docstring'''



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

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

										a__  :	Union[str, Any]		 =			["pixel_values"]
										self.assertListEqual(arg_names[:1]	, a_    )
				def       UpperCAmelCase      (		self    :			Any    )	->	Optional[int]:



							'''simple docstring'''



							a__  :	Dict		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*a_    )
				def       UpperCAmelCase      (		self    :			Tuple    )	->	Dict:



							'''simple docstring'''



							a__  :	List[Any]		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*a_    )
				def       UpperCAmelCase      (		self    :			Any	, a_    :			Any	, a_    :			str	, a_    :			Any    )	->	List[Any]:



							'''simple docstring'''



							np.random.seed(2    )

							a__  :	Dict		 =			int((pt_model.config.image_size // pt_model.config.patch_size) ** 2    )
							a__  :	List[str]		 =			np.random.uniform(size=(self.model_tester.batch_size, num_patches)    )
							a__  :	Optional[Any]		 =			torch.from_numpy(a_    )

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

							super().check_pt_tf_models(a_	, a_	, a_    )
				def       UpperCAmelCase      (		self    :			Dict    )	->	Union[str, Any]:



							'''simple docstring'''



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

							for model_class in self.all_model_classes:
										a__  :	Union[str, Any]		 =			model_class(a_    )
										model.to(a_    )
										model.eval()
										# make random mask reproducible
										torch.manual_seed(2    )
										with torch.no_grad():
													a__  :	Tuple		 =			model(**self._prepare_for_class(a_	, a_    )    )

										a__  :	str		 =			outputs[0].cpu().numpy()
										a__  :	Dict		 =			0

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

													# Make sure we don't have nans
													a__  :	Tuple		 =			after_outputs[0].cpu().numpy()
													a__  :	str		 =			0
													a__  :	Optional[int]		 =			np.amax(np.abs(out_a - out_a    )    )
													self.assertLessEqual(a_	, 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    :			Optional[Any]    )	->	str:



							'''simple docstring'''



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



							'''simple docstring'''



							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    :			Union[str, Any]    )	->	Optional[int]:



							'''simple docstring'''



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



							'''simple docstring'''



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



							'''simple docstring'''



							pass






				@slow
				def       UpperCAmelCase      (		self    :			str    )	->	Any:



							'''simple docstring'''



							for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										a__  :	Optional[int]		 =			ViTMAEModel.from_pretrained(a_    )
										self.assertIsNotNone(a_    )






def      lowercase__ (					)    ->			Optional[int]:



			'''simple docstring'''
			a__  :	Union[str, Any]		 =			Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
			return image





@require_torch
@require_vision
class 					__UpperCAmelCase							(						unittest.TestCase  ):
				@cached_property
				def       UpperCAmelCase      (		self    :			List[Any]    )	->	str:



							'''simple docstring'''



							return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"    ) if is_vision_available() else None






				@slow
				def       UpperCAmelCase      (		self    :			Tuple    )	->	Union[str, Any]:



							'''simple docstring'''



							np.random.seed(2    )

							a__  :	Any		 =			ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"    ).to(a_    )

							a__  :	Union[str, Any]		 =			self.default_image_processor
							a__  :	str		 =			prepare_img()
							a__  :	Tuple		 =			image_processor(images=a_	, return_tensors="pt"    ).to(a_    )

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

							# forward pass
							with torch.no_grad():
										a__  :	Dict		 =			model(**a_	, noise=torch.from_numpy(a_    ).to(device=a_    )    )

							# verify the logits
							a__  :	Dict		 =			torch.Size((1, 1_96, 7_68)    )
							self.assertEqual(outputs.logits.shape	, a_    )

							a__  :	Any		 =			torch.tensor(
							    [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]    )

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






"""simple docstring"""






import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()






def      lowercase__ (					lowerCAmelCase__   :			Union[str, Any]							,			lowerCAmelCase__   :			Dict							,			lowerCAmelCase__   :			List[Any]=None			)    ->			str:



			'''simple docstring'''
			# set parameter of one layer
			assert torch_layer.weight.shape == weight.shape, F"{torch_layer} layer.weight does not match"
			a__  :	Optional[Any]		 =			nn.Parameter(lowerCAmelCase__			)
			if bias is not None:
						assert torch_layer.bias.shape == bias.shape, F"{torch_layer} layer.bias does not match"
						a__  :	str		 =			nn.Parameter(lowerCAmelCase__			)






def      lowercase__ (					lowerCAmelCase__   :			Union[str, Any]							,			lowerCAmelCase__   :			Tuple							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			str:



			'''simple docstring'''
			# set torch weights for 1-to-1 comparison
			a__  :	Tuple		 =			np.asarray(weights[0]			)
			a__  :	Any		 =			np.asarray(weights[1]			)
			a__  :	Optional[Any]		 =			np.asarray(weights[2]			)

			set_param(
			    torch_layer.self_attention.query_key							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.self_attention.value							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.output.dense							,			torch.tensor(lowerCAmelCase__			).view(-1							,			lowerCAmelCase__			).contiguous().transpose(0							,			1			)							,			)






def      lowercase__ (					lowerCAmelCase__   :			List[str]							,			lowerCAmelCase__   :			List[str]							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			str:



			'''simple docstring'''
			# set torch weights for 1-to-1 comparison
			a__  :	int		 =			np.asarray(weights[0]			)
			a__  :	List[str]		 =			np.asarray(weights[1]			)
			a__  :	Optional[int]		 =			np.asarray(weights[2]			)
			a__  :	int		 =			np.asarray(weights[3]			)

			set_param(
			    torch_layer.self_attention.query							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.self_attention.key							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.self_attention.value							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.output.dense							,			torch.tensor(lowerCAmelCase__			).view(-1							,			lowerCAmelCase__			).contiguous().transpose(0							,			1			)							,			)






def      lowercase__ (					lowerCAmelCase__   :			Any							,			lowerCAmelCase__   :			List[Any]							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			Any:



			'''simple docstring'''
			# layernorm 1
			a__  :	int		 =			weights[0][0][0]
			a__  :	Optional[Any]		 =			np.asarray(layer_norm_a[0]			)
			a__  :	Optional[int]		 =			np.asarray(layer_norm_a[1]			)
			set_param(
			    torch_block.attention.layer_norm							,			torch.tensor(lowerCAmelCase__			)							,			torch.tensor(lowerCAmelCase__			)							,			)

			# lsh weights + output
			a__  :	Dict		 =			weights[0][1]
			if len(lowerCAmelCase__			) < 4:
						set_layer_weights_in_torch_lsh(lowerCAmelCase__							,			torch_block.attention							,			lowerCAmelCase__			)
			else:
						set_layer_weights_in_torch_local(lowerCAmelCase__							,			torch_block.attention							,			lowerCAmelCase__			)

			# intermediate weighs
			a__  :	Optional[int]		 =			weights[2][0][1][2]

			# Chunked Feed Forward
			if len(lowerCAmelCase__			) == 4:
						a__  :	List[str]		 =			intermediate_weights[2]

			# layernorm 2
			a__  :	List[Any]		 =			np.asarray(intermediate_weights[0][0]			)
			a__  :	Optional[int]		 =			np.asarray(intermediate_weights[0][1]			)
			set_param(
			    torch_block.feed_forward.layer_norm							,			torch.tensor(lowerCAmelCase__			)							,			torch.tensor(lowerCAmelCase__			)							,			)

			# intermediate dense
			a__  :	Tuple		 =			np.asarray(intermediate_weights[1][0]			)
			a__  :	Union[str, Any]		 =			np.asarray(intermediate_weights[1][1]			)
			set_param(
			    torch_block.feed_forward.dense.dense							,			torch.tensor(lowerCAmelCase__			).transpose(0							,			1			).contiguous()							,			torch.tensor(lowerCAmelCase__			)							,			)

			# intermediate out
			a__  :	Optional[Any]		 =			np.asarray(intermediate_weights[4][0]			)
			a__  :	Dict		 =			np.asarray(intermediate_weights[4][1]			)
			set_param(
			    torch_block.feed_forward.output.dense							,			torch.tensor(lowerCAmelCase__			).transpose(0							,			1			).contiguous()							,			torch.tensor(lowerCAmelCase__			)							,			)






def      lowercase__ (					lowerCAmelCase__   :			int							,			lowerCAmelCase__   :			List[str]							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			Any:



			'''simple docstring'''
			# reformer model
			a__  :	int		 =			torch_model.reformer

			# word embeds
			a__  :	List[Any]		 =			np.asarray(weights[1]			)
			set_param(
			    torch_model_reformer.embeddings.word_embeddings							,			torch.tensor(lowerCAmelCase__			)							,			)

			if isinstance(weights[3]							,			lowerCAmelCase__			):
						a__  :	Dict		 =			torch_model_reformer.embeddings.position_embeddings
						for emb_idx in range(len(position_embeddings.weights			)			):
									a__  :	str		 =			np.asarray(weights[3][emb_idx][0]			)
									assert (
									    position_embeddings.weights[emb_idx].shape == emb_weights.shape
									), F"{position_embeddings[emb_idx]} emb does not match"
									a__  :	List[str]		 =			nn.Parameter(torch.tensor(lowerCAmelCase__			)			)

			a__  :	Optional[Any]		 =			weights[5]
			assert len(torch_model_reformer.encoder.layers			) * 4 == len(
			    lowerCAmelCase__			), "HF and trax model do not have the same number of layers"
			for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers			):
						a__  :	Union[str, Any]		 =			trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
						set_block_weights_in_torch(lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__			)

			# output layer norm
			a__  :	str		 =			np.asarray(weights[7][0]			)
			a__  :	List[Any]		 =			np.asarray(weights[7][1]			)
			set_param(
			    torch_model_reformer.encoder.layer_norm							,			torch.tensor(lowerCAmelCase__			)							,			torch.tensor(lowerCAmelCase__			)							,			)

			# output embeddings
			a__  :	str		 =			np.asarray(weights[9][0]			)
			a__  :	Optional[int]		 =			np.asarray(weights[9][1]			)
			set_param(
			    torch_model.lm_head.decoder							,			torch.tensor(lowerCAmelCase__			).transpose(0							,			1			).contiguous()							,			torch.tensor(lowerCAmelCase__			)							,			)






def      lowercase__ (					lowerCAmelCase__   :			Tuple							,			lowerCAmelCase__   :			List[Any]							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			List[str]:



			'''simple docstring'''
			# Initialise PyTorch model
			a__  :	List[Any]		 =			ReformerConfig.from_json_file(lowerCAmelCase__			)
			print(F"Building PyTorch model from configuration: {config}"			)
			a__  :	Optional[Any]		 =			ReformerModelWithLMHead(lowerCAmelCase__			)

			with open(lowerCAmelCase__							,			"rb"			) as f:
						a__  :	List[str]		 =			pickle.load(lowerCAmelCase__			)["weights"]

			set_model_weights_in_torch(lowerCAmelCase__							,			lowerCAmelCase__							,			config.hidden_size			)

			# Save pytorch-model
			print(F"Save PyTorch model to {pytorch_dump_path}"			)
			torch.save(model.state_dict()							,			lowerCAmelCase__			)


if __name__ == "__main__":
		__UpperCAmelCase						=			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--trax_model_pkl_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
		)
		parser.add_argument(
		    '''--config_file''',
		    default=None,
		    type=str,
		    required=True,
		    help=(
		        '''The config json file corresponding to the pre-trained Reformer model. \n'''
		        '''This specifies the model architecture.'''
		    ),
		)
		parser.add_argument(
		    '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
		)
		__UpperCAmelCase						=			parser.parse_args()
		convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)