"""simple docstring"""




import inspect
import unittest

from transformers import ViTHybridConfig
from transformers.testing_utils import require_accelerate, 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, _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 ViTHybridForImageClassification, ViTHybridImageProcessor, ViTHybridModel
				from transformers.models.vit_hybrid.modeling_vit_hybrid import VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
				from PIL import Image
class 			UpperCamelCase__:





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase=13						,__UpperCAmelCase=64						,__UpperCAmelCase=2						,__UpperCAmelCase=3						,__UpperCAmelCase=True						,__UpperCAmelCase=True						,__UpperCAmelCase=32						,__UpperCAmelCase=5						,__UpperCAmelCase=4						,__UpperCAmelCase=37						,__UpperCAmelCase="gelu"						,__UpperCAmelCase=0.1						,__UpperCAmelCase=0.1						,__UpperCAmelCase=10						,__UpperCAmelCase=0.0_2						,__UpperCAmelCase=[1, 16, 4, 4]						,__UpperCAmelCase=None						,)    -> List[Any]:
													A__		   =							parent
													A__		   =							batch_size
													A__		   =							image_size
													A__		   =							patch_size
													A__		   =							num_channels
													A__		   =							is_training
													A__		   =							use_labels
													A__		   =							hidden_size
													A__		   =							num_hidden_layers
													A__		   =							num_attention_heads
													A__		   =							intermediate_size
													A__		   =							hidden_act
													A__		   =							hidden_dropout_prob
													A__		   =							attention_probs_dropout_prob
													A__		   =							type_sequence_label_size
													A__		   =							initializer_range
													A__		   =							scope
													A__		   =							backbone_featmap_shape

													# in ViT hybrid, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
													# the number of patches is based on the feature map of the backbone, which by default uses an output stride
													# of 32, which means that the feature map has a spatial resolution of 1/32 of the input image size
													A__		   =							(self.image_size // 32) ** 2
													A__		   =							num_patches + 1





						def  snake_case__		(				self   )    -> Optional[int]:
													A__		   =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

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

													A__		   =							self.get_config()

													return config, pixel_values, labels





						def  snake_case__		(				self   )    -> List[Any]:
													A__		   =							{
													    'global_padding': 'same',
													    'layer_type': 'bottleneck',
													    'depths': [3, 4, 9],
													    'out_features': ['stage1', 'stage2', 'stage3'],
													    'embedding_dynamic_padding': True,
													    'hidden_sizes': [4, 8, 16, 32],
													    'num_groups': 2,
													}

													return ViTHybridConfig(
													    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=__UpperCAmelCase						,initializer_range=self.initializer_range						,backbone_featmap_shape=self.backbone_featmap_shape						,backbone_config=__UpperCAmelCase						,)





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )    -> int:
													A__		   =							ViTHybridModel(config=__UpperCAmelCase   )
													model.to(__UpperCAmelCase   )
													model.eval()
													A__		   =							model(__UpperCAmelCase   )
													self.parent.assertEqual(result.last_hidden_state.shape						,(self.batch_size, self.seq_length, self.hidden_size)   )





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )    -> int:
													A__		   =							self.type_sequence_label_size
													A__		   =							ViTHybridForImageClassification(__UpperCAmelCase   )
													model.to(__UpperCAmelCase   )
													model.eval()
													A__		   =							model(__UpperCAmelCase						,labels=__UpperCAmelCase   )
													self.parent.assertEqual(result.logits.shape						,(self.batch_size, self.type_sequence_label_size)   )





						def  snake_case__		(				self   )    -> Optional[Any]:
													A__		   =							self.prepare_config_and_inputs()
													A__   ,     A__   ,     A__		   =							config_and_inputs
													A__		   =							{'pixel_values': pixel_values}
													return config, inputs_dict






@require_torch
class 			UpperCamelCase__(	__A						,  __A						,  unittest.TestCase       ):
						lowerCAmelCase__     :  Tuple  							=						(ViTHybridModel, ViTHybridForImageClassification) if is_torch_available() else ()
						lowerCAmelCase__     :  Any  							=						(
						    {'feature-extraction': ViTHybridModel, 'image-classification': ViTHybridForImageClassification}
						    if is_torch_available()
						    else {}
						)
						lowerCAmelCase__     :  Any  							=						False
						lowerCAmelCase__     :  Any  							=						False
						lowerCAmelCase__     :  Optional[int]  							=						False





						def  snake_case__		(				self   )    -> Tuple:
													A__		   =							ViTHybridModelTester(self   )
													A__		   =							ConfigTester(self						,config_class=__UpperCAmelCase						,has_text_modality=__UpperCAmelCase						,hidden_size=37   )





						def  snake_case__		(				self   )    -> Tuple:
													self.config_tester.run_common_tests()





						@unittest.skip(reason='ViT does not use inputs_embeds'   )
						def  snake_case__		(				self   )    -> int:
													pass





						def  snake_case__		(				self   )    -> Dict:
													A__   ,     A__		   =							self.model_tester.prepare_config_and_inputs_for_common()

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





						def  snake_case__		(				self   )    -> Optional[Any]:
													A__   ,     A__		   =							self.model_tester.prepare_config_and_inputs_for_common()

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

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





						def  snake_case__		(				self   )    -> Union[str, Any]:
													A__		   =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*__UpperCAmelCase   )





						def  snake_case__		(				self   )    -> Optional[int]:
													A__		   =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase   )





						def  snake_case__		(				self   )    -> Any:
													A__   ,     A__		   =							self.model_tester.prepare_config_and_inputs_for_common()

													A__		   =							_config_zero_init(__UpperCAmelCase   )
													for model_class in self.all_model_classes:
																				A__		   =							model_class(config=__UpperCAmelCase   )
																				# Skip the check for the backbone
																				for name, module in model.named_modules():
																											if module.__class__.__name__ == "ViTHybridPatchEmbeddings":
																																		A__		   =							[f'''{name}.{key}''' for key in module.state_dict().keys()]
																																		break

																				for name, param in model.named_parameters():
																											if param.requires_grad:
																																		if name in backbone_params:
																																									continue
																																		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'''						,)





						@slow
						def  snake_case__		(				self   )    -> Optional[Any]:
													for model_name in VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				A__		   =							ViTHybridModel.from_pretrained(__UpperCAmelCase   )
																				self.assertIsNotNone(__UpperCAmelCase   )






def      UpperCAmelCase    (		):






							"""simple docstring"""





							A__		   =							Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png' )
							return image


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





						@cached_property
						def  snake_case__		(				self   )    -> List[Any]:
													return (
													    ViTHybridImageProcessor.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]   )
													    if is_vision_available()
													    else None
													)





						@slow
						def  snake_case__		(				self   )    -> Optional[Any]:
													A__		   =							ViTHybridForImageClassification.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]   ).to(
													    __UpperCAmelCase   )

													A__		   =							self.default_image_processor
													A__		   =							prepare_img()
													A__		   =							image_processor(images=__UpperCAmelCase						,return_tensors='pt'   ).to(__UpperCAmelCase   )

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

													# verify the logits
													A__		   =							torch.Size((1, 10_00)   )
													self.assertEqual(outputs.logits.shape						,__UpperCAmelCase   )

													A__		   =							torch.tensor([-1.9_0_9_0, -0.4_9_9_3, -0.2_3_8_9]   ).to(__UpperCAmelCase   )

													self.assertTrue(torch.allclose(outputs.logits[0, :3]						,__UpperCAmelCase						,atol=1e-4   )   )





						@slow
						@require_accelerate
						def  snake_case__		(				self   )    -> List[Any]:
													A__		   =							ViTHybridImageProcessor.from_pretrained('google/vit-hybrid-base-bit-384'   )
													A__		   =							ViTHybridForImageClassification.from_pretrained('google/vit-hybrid-base-bit-384'						,device_map='auto'   )

													A__		   =							prepare_img()

													A__		   =							image_processor(images=__UpperCAmelCase						,return_tensors='pt'   )
													A__		   =							model(**__UpperCAmelCase   )
													A__		   =							outputs.logits
													# model predicts one of the 1000 ImageNet classes
													A__		   =							logits.argmax(-1   ).item()

													self.assertTrue(model.config.idalabel[predicted_class_idx]						,'tabby, tabby cat'   )






"""simple docstring"""




import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
				from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
				import torch_xla.core.xla_model as xm






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							if is_torch_version('<'				,			'2.0.0' ) or not hasattr(UpperCamelCase__				,			'_dynamo' ):
														return False
							return isinstance(UpperCamelCase__				,			torch._dynamo.eval_frame.OptimizedModule )






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ = True ):






							"""simple docstring"""





							A__		   =							(torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

							A__		   =							is_compiled_module(UpperCamelCase__ )
							if is_compiled:
														A__		   =							model
														A__		   =							model._orig_mod

							if is_deepspeed_available():
														options += (DeepSpeedEngine,)

							while isinstance(UpperCamelCase__				,			UpperCamelCase__ ):
														A__		   =							model.module

							if not keep_fpaa_wrapper:
														A__		   =							getattr(UpperCamelCase__				,			'forward' )
														A__		   =							model.__dict__.pop('_original_forward'				,			UpperCamelCase__ )
														if original_forward is not None:
																					while hasattr(UpperCamelCase__				,			'__wrapped__' ):
																												A__		   =							forward.__wrapped__
																												if forward == original_forward:
																																			break
																					A__		   =							forward
														if getattr(UpperCamelCase__				,			'_converted_to_transformer_engine'				,			UpperCamelCase__ ):
																					convert_model(UpperCamelCase__				,			to_transformer_engine=UpperCamelCase__ )

							if is_compiled:
														A__		   =							model
														A__		   =							compiled_model

							return model






def      UpperCAmelCase    (		):






							"""simple docstring"""





							PartialState().wait_for_everyone()






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							if PartialState().distributed_type == DistributedType.TPU:
														xm.save(UpperCamelCase__				,			UpperCamelCase__ )
							elif PartialState().local_process_index == 0:
														torch.save(UpperCamelCase__				,			UpperCamelCase__ )






@contextmanager
def      UpperCAmelCase    (		**UpperCamelCase__ ):






							"""simple docstring"""





							for key, value in kwargs.items():
														A__		   =							str(UpperCamelCase__ )

							yield

							for key in kwargs:
														if key.upper() in os.environ:
																					del os.environ[key.upper()]






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							if not hasattr(UpperCamelCase__				,			'__qualname__' ) and not hasattr(UpperCamelCase__				,			'__name__' ):
														A__		   =							getattr(UpperCamelCase__				,			'__class__'				,			UpperCamelCase__ )
							if hasattr(UpperCamelCase__				,			'__qualname__' ):
														return obj.__qualname__
							if hasattr(UpperCamelCase__				,			'__name__' ):
														return obj.__name__
							return str(UpperCamelCase__ )






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							for key, value in source.items():
														if isinstance(UpperCamelCase__				,			UpperCamelCase__ ):
																					A__		   =							destination.setdefault(UpperCamelCase__				,			{} )
																					merge_dicts(UpperCamelCase__				,			UpperCamelCase__ )
														else:
																					A__		   =							value

							return destination






def      UpperCAmelCase    (		UpperCamelCase__ = None ):






							"""simple docstring"""





							if port is None:
														A__		   =							29_500
							with socket.socket(socket.AF_INET				,			socket.SOCK_STREAM ) as s:
														return s.connect_ex(('localhost', port) ) == 0





