
import collections
import inspect
import unittest

from transformers import SwinConfig
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_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import SwinBackbone, SwinForImageClassification, SwinForMaskedImageModeling, SwinModel
							from transformers.models.swin.modeling_swin import SWIN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor







class 							lowercase						:




							def __init__(							self  ,A__  ,A__=1_3  ,A__=3_2  ,A__=2  ,A__=3  ,A__=1_6  ,A__=[1, 2, 1]  ,A__=[2, 2, 4]  ,A__=2  ,A__=2.0  ,A__=True  ,A__=0.0  ,A__=0.0  ,A__=0.1  ,A__="gelu"  ,A__=False  ,A__=True  ,A__=0.02  ,A__=1E-5  ,A__=True  ,A__=None  ,A__=True  ,A__=1_0  ,A__=8  ,A__=["stage1", "stage2"]  ,A__=[1, 2]  ,):
												lowercase		      =		parent
												lowercase		      =		batch_size
												lowercase		      =		image_size
												lowercase		      =		patch_size
												lowercase		      =		num_channels
												lowercase		      =		embed_dim
												lowercase		      =		depths
												lowercase		      =		num_heads
												lowercase		      =		window_size
												lowercase		      =		mlp_ratio
												lowercase		      =		qkv_bias
												lowercase		      =		hidden_dropout_prob
												lowercase		      =		attention_probs_dropout_prob
												lowercase		      =		drop_path_rate
												lowercase		      =		hidden_act
												lowercase		      =		use_absolute_embeddings
												lowercase		      =		patch_norm
												lowercase		      =		layer_norm_eps
												lowercase		      =		initializer_range
												lowercase		      =		is_training
												lowercase		      =		scope
												lowercase		      =		use_labels
												lowercase		      =		type_sequence_label_size
												lowercase		      =		encoder_stride
												lowercase		      =		out_features
												lowercase		      =		out_indices




							def 					A__ (							self):
												lowercase		      =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

												lowercase		      =		self.get_config()

												return config, pixel_values, labels




							def 					A__ (							self):
												return SwinConfig(
												    image_size=self.image_size  ,patch_size=self.patch_size  ,num_channels=self.num_channels  ,embed_dim=self.embed_dim  ,depths=self.depths  ,num_heads=self.num_heads  ,window_size=self.window_size  ,mlp_ratio=self.mlp_ratio  ,qkv_bias=self.qkv_bias  ,hidden_dropout_prob=self.hidden_dropout_prob  ,attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,drop_path_rate=self.drop_path_rate  ,hidden_act=self.hidden_act  ,use_absolute_embeddings=self.use_absolute_embeddings  ,path_norm=self.patch_norm  ,layer_norm_eps=self.layer_norm_eps  ,initializer_range=self.initializer_range  ,encoder_stride=self.encoder_stride  ,out_features=self.out_features  ,out_indices=self.out_indices  ,)




							def 					A__ (							self  ,A__  ,A__  ,A__):
												lowercase		      =		SwinModel(config=A__)
												model.to(A__)
												model.eval()
												lowercase		      =		model(A__)

												lowercase		      =		((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths) - 1))
												lowercase		      =		int(config.embed_dim * 2 ** (len(config.depths) - 1))

												self.parent.assertEqual(result.last_hidden_state.shape  ,(self.batch_size, expected_seq_len, expected_dim))




							def 					A__ (							self  ,A__  ,A__  ,A__):
												lowercase		      =		SwinBackbone(config=A__)
												model.to(A__)
												model.eval()
												lowercase		      =		model(A__)

												# verify hidden states
												self.parent.assertEqual(len(result.feature_maps)  ,len(config.out_features))
												self.parent.assertListEqual(list(result.feature_maps[0].shape)  ,[self.batch_size, model.channels[0], 1_6, 1_6])

												# verify channels
												self.parent.assertEqual(len(model.channels)  ,len(config.out_features))

												# verify backbone works with out_features=None
												lowercase		      =		None
												lowercase		      =		SwinBackbone(config=A__)
												model.to(A__)
												model.eval()
												lowercase		      =		model(A__)

												# verify feature maps
												self.parent.assertEqual(len(result.feature_maps)  ,1)
												self.parent.assertListEqual(list(result.feature_maps[0].shape)  ,[self.batch_size, model.channels[-1], 4, 4])

												# verify channels
												self.parent.assertEqual(len(model.channels)  ,1)




							def 					A__ (							self  ,A__  ,A__  ,A__):
												lowercase		      =		SwinForMaskedImageModeling(config=A__)
												model.to(A__)
												model.eval()
												lowercase		      =		model(A__)
												self.parent.assertEqual(
												    result.logits.shape  ,(self.batch_size, self.num_channels, self.image_size, self.image_size))

												# test greyscale images
												lowercase		      =		1
												lowercase		      =		SwinForMaskedImageModeling(A__)
												model.to(A__)
												model.eval()

												lowercase		      =		floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
												lowercase		      =		model(A__)
												self.parent.assertEqual(result.logits.shape  ,(self.batch_size, 1, self.image_size, self.image_size))




							def 					A__ (							self  ,A__  ,A__  ,A__):
												lowercase		      =		self.type_sequence_label_size
												lowercase		      =		SwinForImageClassification(A__)
												model.to(A__)
												model.eval()
												lowercase		      =		model(A__  ,labels=A__)
												self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.type_sequence_label_size))

												# test greyscale images
												lowercase		      =		1
												lowercase		      =		SwinForImageClassification(A__)
												model.to(A__)
												model.eval()

												lowercase		      =		floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
												lowercase		      =		model(A__)
												self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.type_sequence_label_size))





							def 					A__ (							self):
												lowercase		      =		self.prepare_config_and_inputs()
												(
												    (
												    lowercase
												)      ,		(
												    lowercase
												)      ,		(
												    lowercase
												)      ,		
												)		      =		config_and_inputs
												lowercase		      =		{'''pixel_values''': pixel_values}
												return config, inputs_dict







@require_torch
class 							lowercase						(   SCREAMING_SNAKE_CASE__			,    SCREAMING_SNAKE_CASE__			,    unittest.TestCase						):
							lowercase_			:			Optional[Any]			   =(
							    (
							        SwinModel,
							        SwinBackbone,
							        SwinForImageClassification,
							        SwinForMaskedImageModeling,
							    )
							    if is_torch_available()
							    else ()
							)
							lowercase_			:			List[str]			   =(
							    {'''feature-extraction''': SwinModel, '''image-classification''': SwinForImageClassification}
							    if is_torch_available()
							    else {}
							)
							lowercase_			:			Dict			   =True

							lowercase_			:			Any			   =False
							lowercase_			:			Any			   =False
							lowercase_			:			List[Any]			   =False




							def 					A__ (							self):
												lowercase		      =		SwinModelTester(self)
												lowercase		      =		ConfigTester(self  ,config_class=A__  ,embed_dim=3_7)




							def 					A__ (							self):
												self.create_and_test_config_common_properties()
												self.config_tester.create_and_test_config_to_json_string()
												self.config_tester.create_and_test_config_to_json_file()
												self.config_tester.create_and_test_config_from_and_save_pretrained()
												self.config_tester.create_and_test_config_with_num_labels()
												self.config_tester.check_config_can_be_init_without_params()
												self.config_tester.check_config_arguments_init()




							def 					A__ (							self):
												return




							def 					A__ (							self):
												lowercase		      =		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*A__)




							@unittest.skip(reason='''Got `CUDA error: misaligned address` with PyTorch 2.0.0.''')
							def 					A__ (							self):
												pass




							def 					A__ (							self):
												super().test_training_gradient_checkpointing()




							def 					A__ (							self):
												lowercase		      =		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_backbone(*A__)




							def 					A__ (							self):
												lowercase		      =		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_image_modeling(*A__)




							def 					A__ (							self):
												lowercase		      =		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*A__)




							@unittest.skip(reason='''Swin does not use inputs_embeds''')
							def 					A__ (							self):
												pass




							@unittest.skip(reason='''Swin Transformer does not use feedforward chunking''')
							def 					A__ (							self):
												pass




							def 					A__ (							self):
												lowercase      ,		lowercase		      =		self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																	lowercase		      =		model_class(A__)
																	self.assertIsInstance(model.get_input_embeddings()  ,(nn.Module))
																	lowercase		      =		model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(A__  ,nn.Linear))




							def 					A__ (							self):
												lowercase      ,		lowercase		      =		self.model_tester.prepare_config_and_inputs_for_common()

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

																	lowercase		      =		['''pixel_values''']
																	self.assertListEqual(arg_names[:1]  ,A__)




							def 					A__ (							self):
												lowercase      ,		lowercase		      =		self.model_tester.prepare_config_and_inputs_for_common()
												lowercase		      =		True

												for model_class in self.all_model_classes:
																	lowercase		      =		True
																	lowercase		      =		False
																	lowercase		      =		True
																	lowercase		      =		model_class(A__)
																	model.to(A__)
																	model.eval()
																	with torch.no_grad():
																						lowercase		      =		model(**self._prepare_for_class(A__  ,A__))
																	lowercase		      =		outputs.attentions
																	lowercase		      =		len(self.model_tester.depths)
																	self.assertEqual(len(A__)  ,A__)

																	# check that output_attentions also work using config
																	del inputs_dict["output_attentions"]
																	lowercase		      =		True
																	lowercase		      =		config.window_size**2
																	lowercase		      =		model_class(A__)
																	model.to(A__)
																	model.eval()
																	with torch.no_grad():
																						lowercase		      =		model(**self._prepare_for_class(A__  ,A__))
																	lowercase		      =		outputs.attentions
																	self.assertEqual(len(A__)  ,A__)

																	self.assertListEqual(
																	    list(attentions[0].shape[-3:])  ,[self.model_tester.num_heads[0], window_size_squared, window_size_squared]  ,)
																	lowercase		      =		len(A__)

																	# Check attention is always last and order is fine
																	lowercase		      =		True
																	lowercase		      =		True
																	lowercase		      =		model_class(A__)
																	model.to(A__)
																	model.eval()
																	with torch.no_grad():
																						lowercase		      =		model(**self._prepare_for_class(A__  ,A__))

																	# also another +1 for reshaped_hidden_states
																	lowercase		      =		1 if model_class.__name__ == '''SwinBackbone''' else 2
																	self.assertEqual(out_len + added_hidden_states  ,len(A__))

																	lowercase		      =		outputs.attentions

																	self.assertEqual(len(A__)  ,A__)

																	self.assertListEqual(
																	    list(self_attentions[0].shape[-3:])  ,[self.model_tester.num_heads[0], window_size_squared, window_size_squared]  ,)




							def 					A__ (							self  ,A__  ,A__  ,A__  ,A__):
												lowercase		      =		model_class(A__)
												model.to(A__)
												model.eval()

												with torch.no_grad():
																	lowercase		      =		model(**self._prepare_for_class(A__  ,A__))

												lowercase		      =		outputs.hidden_states

												lowercase		      =		getattr(
												    self.model_tester  ,'''expected_num_hidden_layers'''  ,len(self.model_tester.depths) + 1)
												self.assertEqual(len(A__)  ,A__)

												# Swin has a different seq_length
												lowercase		      =		(
												    config.patch_size
												    if isinstance(config.patch_size  ,collections.abc.Iterable)
												    else (config.patch_size, config.patch_size)
												)

												lowercase		      =		(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

												self.assertListEqual(
												    list(hidden_states[0].shape[-2:])  ,[num_patches, self.model_tester.embed_dim]  ,)

												if not model_class.__name__ == "SwinBackbone":
																	lowercase		      =		outputs.reshaped_hidden_states
																	self.assertEqual(len(A__)  ,A__)

																	lowercase      ,		lowercase      ,		lowercase      ,		lowercase		      =		reshaped_hidden_states[0].shape
																	lowercase		      =		(
																	    reshaped_hidden_states[0].view(A__  ,A__  ,height * width).permute(0  ,2  ,1)
																	)
																	self.assertListEqual(
																	    list(reshaped_hidden_states.shape[-2:])  ,[num_patches, self.model_tester.embed_dim]  ,)




							def 					A__ (							self):
												lowercase      ,		lowercase		      =		self.model_tester.prepare_config_and_inputs_for_common()

												lowercase		      =		(
												    self.model_tester.image_size
												    if isinstance(self.model_tester.image_size  ,collections.abc.Iterable)
												    else (self.model_tester.image_size, self.model_tester.image_size)
												)

												for model_class in self.all_model_classes:
																	lowercase		      =		True
																	self.check_hidden_states_output(A__  ,A__  ,A__  ,A__)

																	# check that output_hidden_states also work using config
																	del inputs_dict["output_hidden_states"]
																	lowercase		      =		True

																	self.check_hidden_states_output(A__  ,A__  ,A__  ,A__)




							def 					A__ (							self):
												lowercase      ,		lowercase		      =		self.model_tester.prepare_config_and_inputs_for_common()
												lowercase		      =		3

												lowercase		      =		(
												    self.model_tester.image_size
												    if isinstance(self.model_tester.image_size  ,collections.abc.Iterable)
												    else (self.model_tester.image_size, self.model_tester.image_size)
												)
												lowercase		      =		(
												    config.patch_size
												    if isinstance(config.patch_size  ,collections.abc.Iterable)
												    else (config.patch_size, config.patch_size)
												)

												lowercase		      =		image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
												lowercase		      =		image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

												for model_class in self.all_model_classes:
																	lowercase		      =		True
																	self.check_hidden_states_output(A__  ,A__  ,A__  ,(padded_height, padded_width))

																	# check that output_hidden_states also work using config
																	del inputs_dict["output_hidden_states"]
																	lowercase		      =		True
																	self.check_hidden_states_output(A__  ,A__  ,A__  ,(padded_height, padded_width))




							@slow
							def 					A__ (							self):
												for model_name in SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase		      =		SwinModel.from_pretrained(A__)
																	self.assertIsNotNone(A__)





							def 					A__ (							self):
												lowercase      ,		lowercase		      =		self.model_tester.prepare_config_and_inputs_for_common()

												lowercase		      =		_config_zero_init(A__)
												for model_class in self.all_model_classes:
																	lowercase		      =		model_class(config=A__)
																	for name, param in model.named_parameters():
																						if "embeddings" not in name and param.requires_grad:
																											self.assertIn(
																											    ((param.data.mean() * 1E9).round() / 1E9).item()  ,[0.0, 1.0]  ,msg=f'Parameter {name} of model {model_class} seems not properly initialized'  ,)







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




							@cached_property
							def 					A__ (							self):
												return (
												    AutoImageProcessor.from_pretrained('''microsoft/swin-tiny-patch4-window7-224''')
												    if is_vision_available()
												    else None
												)





							@slow
							def 					A__ (							self):
												lowercase		      =		SwinForImageClassification.from_pretrained('''microsoft/swin-tiny-patch4-window7-224''').to(A__)
												lowercase		      =		self.default_image_processor

												lowercase		      =		Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
												lowercase		      =		image_processor(images=A__  ,return_tensors='''pt''').to(A__)

												# forward pass
												with torch.no_grad():
																	lowercase		      =		model(**A__)

												# verify the logits
												lowercase		      =		torch.Size((1, 1_0_0_0))
												self.assertEqual(outputs.logits.shape  ,A__)
												lowercase		      =		torch.tensor([-0.0948, -0.6454, -0.0921]).to(A__)
												self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,A__  ,atol=1E-4))







@require_torch
class 							lowercase						(   unittest.TestCase			,    SCREAMING_SNAKE_CASE__						):
							lowercase_			:			List[Any]			   =(SwinBackbone,) if is_torch_available() else ()
							lowercase_			:			Optional[int]			   =SwinConfig




							def 					A__ (							self):
												lowercase		      =		SwinModelTester(self)






import torch

from accelerate import Accelerator, DistributedType







class 							lowercase						:




							def __enter__(							self):
												if self.enabled:
																	lowercase		      =		self.model.no_sync()
																	self.model_sync_obj.__enter__()

												return self




							def __exit__(							self  ,A__  ,A__  ,A__):
												if self.enabled:
																	# Average all models on exit
																	self._sync_and_avg_model_params()
																	self.model_sync_obj.__exit__(A__  ,A__  ,A__)




							def __init__(							self  ,A__  ,A__  ,A__  ,A__ = True):
												if accelerator.distributed_type not in [
												    DistributedType.NO,
												    DistributedType.MULTI_CPU,
												    DistributedType.MULTI_GPU,
												]:
																	raise NotImplementedError('''LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)''')
												lowercase		      =		enabled and accelerator.distributed_type != DistributedType.NO
												lowercase		      =		0
												if self.enabled:
																	lowercase		      =		accelerator
																	lowercase		      =		model
																	lowercase		      =		local_sgd_steps




							def 					A__ (							self):
												self.num_steps += 1
												if not self.enabled:
																	return

												if self.num_steps % self.local_sgd_steps == 0:
																	self._sync_and_avg_model_params()





							def 					A__ (							self):
												self.accelerator.wait_for_everyone()
												with self.accelerator.autocast():
																	for param in self.model.parameters():
																						lowercase		      =		self.accelerator.reduce(param.data  ,reduction='''mean''')




