"""simple docstring"""




import inspect
import unittest
from typing import List

import numpy as np

from transformers import EfficientFormerConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

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 (
				    TFEfficientFormerForImageClassification,
				    TFEfficientFormerForImageClassificationWithTeacher,
				    TFEfficientFormerModel,
				)
				from transformers.models.efficientformer.modeling_tf_efficientformer import (
				    TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
				)


if is_vision_available():
				from PIL import Image

				from transformers import EfficientFormerImageProcessor
class 			UpperCamelCase__:





						def __init__(				self						,__UpperCAmelCase						,__UpperCAmelCase = 13						,__UpperCAmelCase = 64						,__UpperCAmelCase = 2						,__UpperCAmelCase = 3						,__UpperCAmelCase = 3						,__UpperCAmelCase = True						,__UpperCAmelCase = True						,__UpperCAmelCase = 1_28						,__UpperCAmelCase=[16, 32, 64, 1_28]						,__UpperCAmelCase = 7						,__UpperCAmelCase = 4						,__UpperCAmelCase = 37						,__UpperCAmelCase = "gelu"						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 10						,__UpperCAmelCase = 0.0_2						,__UpperCAmelCase = 2						,__UpperCAmelCase = 1						,__UpperCAmelCase = 1_28						,__UpperCAmelCase = [2, 2, 2, 2]						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,)    -> List[str]:
													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__		   =							encoder_stride
													A__		   =							num_attention_outputs
													A__		   =							embed_dim
													A__		   =							embed_dim + 1
													A__		   =							resolution
													A__		   =							depths
													A__		   =							hidden_sizes
													A__		   =							dim
													A__		   =							mlp_expansion_ratio





						def  snake_case__		(				self   )    -> List[str]:
													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]:
													return EfficientFormerConfig(
													    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						,encoder_stride=self.encoder_stride						,resolution=self.resolution						,depths=self.depths						,hidden_sizes=self.hidden_sizes						,dim=self.dim						,mlp_expansion_ratio=self.mlp_expansion_ratio						,)





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )    -> Dict:
													A__		   =							TFEfficientFormerModel(config=__UpperCAmelCase   )
													A__		   =							model(__UpperCAmelCase						,training=__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   )    -> Tuple:
													A__		   =							self.type_sequence_label_size
													A__		   =							TFEfficientFormerForImageClassification(__UpperCAmelCase   )
													A__		   =							model(__UpperCAmelCase						,labels=__UpperCAmelCase						,training=__UpperCAmelCase   )
													self.parent.assertEqual(result.logits.shape						,(self.batch_size, self.type_sequence_label_size)   )

													# test greyscale images
													A__		   =							1
													A__		   =							TFEfficientFormerForImageClassification(__UpperCAmelCase   )

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





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






@require_tf
class 			UpperCamelCase__(	__A						,  __A						,  unittest.TestCase       ):
						lowerCAmelCase__     :  List[str]  							=						(
						    (
						        TFEfficientFormerModel,
						        TFEfficientFormerForImageClassificationWithTeacher,
						        TFEfficientFormerForImageClassification,
						    )
						    if is_tf_available()
						    else ()
						)
						lowerCAmelCase__     :  Union[str, Any]  							=						(
						    {
						        'feature-extraction': TFEfficientFormerModel,
						        'image-classification': (
						            TFEfficientFormerForImageClassification,
						            TFEfficientFormerForImageClassificationWithTeacher,
						        ),
						    }
						    if is_tf_available()
						    else {}
						)

						lowerCAmelCase__     :  str  							=						False

						lowerCAmelCase__     :  Dict  							=						False
						lowerCAmelCase__     :  Dict  							=						False
						lowerCAmelCase__     :  Optional[int]  							=						False
						lowerCAmelCase__     :  Optional[Any]  							=						False





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





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





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





						@unittest.skip(reason='EfficientFormer does not support input and output embeddings'   )
						def  snake_case__		(				self   )    -> List[str]:
													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   )
																				A__		   =							inspect.signature(model.call   )
																				# 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]:
													def check_hidden_states_output(__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   ):
																				A__		   =							model_class(__UpperCAmelCase   )

																				A__		   =							model(**self._prepare_for_class(__UpperCAmelCase						,__UpperCAmelCase   )						,training=__UpperCAmelCase   )
																				A__		   =							outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																				A__		   =							getattr(
																				    self.model_tester						,'expected_num_hidden_layers'						,self.model_tester.num_hidden_layers + 1   )
																				self.assertEqual(len(__UpperCAmelCase   )						,__UpperCAmelCase   )

																				if hasattr(self.model_tester						,'encoder_seq_length'   ):
																											A__		   =							self.model_tester.encoder_seq_length
																											if hasattr(self.model_tester						,'chunk_length'   ) and self.model_tester.chunk_length > 1:
																																		A__		   =							seq_length * self.model_tester.chunk_length
																				else:
																											A__		   =							self.model_tester.seq_length

																				self.assertListEqual(
																				    list(hidden_states[-1].shape[-2:]   )						,[seq_length, self.model_tester.hidden_size]						,)

																				if config.is_encoder_decoder:
																											A__		   =							outputs.decoder_hidden_states

																											self.asseretIsInstance(__UpperCAmelCase						,(list, tuple)   )
																											self.assertEqual(len(__UpperCAmelCase   )						,__UpperCAmelCase   )
																											A__		   =							getattr(self.model_tester						,'seq_length'						,__UpperCAmelCase   )
																											A__		   =							getattr(self.model_tester						,'decoder_seq_length'						,__UpperCAmelCase   )

																											self.assertListEqual(
																											    list(hidden_states[-1].shape[-2:]   )						,[decoder_seq_length, self.model_tester.hidden_size]						,)

													A__   ,     A__		   =							self.model_tester.prepare_config_and_inputs_for_common()

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

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

																				check_hidden_states_output(__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase   )





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase						,__UpperCAmelCase=False   )    -> List[Any]:
													A__		   =							super()._prepare_for_class(__UpperCAmelCase						,__UpperCAmelCase						,return_labels=__UpperCAmelCase   )

													if return_labels:
																				if model_class.__name__ == "TFEfficientFormerForImageClassificationWithTeacher":
																											del inputs_dict["labels"]

													return inputs_dict





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





						@unittest.skip(reason='EfficientFormer does not implement masked image modeling yet'   )
						def  snake_case__		(				self   )    -> Dict:
													A__		   =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_image_modeling(*__UpperCAmelCase   )





						def  snake_case__		(				self   )    -> str:
													A__		   =							self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase   )





						@slow
						def  snake_case__		(				self   )    -> Union[str, Any]:
													for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				A__		   =							TFEfficientFormerModel.from_pretrained(__UpperCAmelCase   )
																				self.assertIsNotNone(__UpperCAmelCase   )





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

													A__		   =							getattr(self.model_tester						,'seq_length'						,__UpperCAmelCase   )
													A__		   =							getattr(self.model_tester						,'encoder_seq_length'						,__UpperCAmelCase   )
													A__		   =							getattr(self.model_tester						,'key_length'						,__UpperCAmelCase   )
													A__		   =							getattr(self.model_tester						,'chunk_length'						,__UpperCAmelCase   )

													if chunk_length is not None and hasattr(self.model_tester						,'num_hashes'   ):
																				A__		   =							encoder_seq_length * self.model_tester.num_hashes

													for model_class in self.all_model_classes:
																				A__		   =							True
																				A__		   =							False
																				A__		   =							True
																				A__		   =							model_class(__UpperCAmelCase   )

																				A__		   =							model(**self._prepare_for_class(__UpperCAmelCase						,__UpperCAmelCase   )						,training=__UpperCAmelCase   )
																				A__		   =							outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
																				self.assertEqual(len(__UpperCAmelCase   )						,self.model_tester.num_attention_outputs   )

																				# check that output_attentions also work using config
																				del inputs_dict["output_attentions"]
																				A__		   =							True
																				A__		   =							model_class(__UpperCAmelCase   )
																				A__		   =							model(**self._prepare_for_class(__UpperCAmelCase						,__UpperCAmelCase   )						,training=__UpperCAmelCase   )

																				A__		   =							outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
																				self.assertEqual(len(__UpperCAmelCase   )						,self.model_tester.num_attention_outputs   )

																				if chunk_length is not None:
																											self.assertListEqual(
																											    list(attentions[0].shape[-4:]   )						,[self.model_tester.num_attention_heads, encoder_seq_length, chunk_length, encoder_key_length]						,)
																				else:
																											self.assertListEqual(
																											    list(attentions[0].shape[-3:]   )						,[self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]						,)





						def  snake_case__		(				self   )    -> List[Any]:
													# We use a simplified version of this test for EfficientFormer because it requires training=False
													# and Keras refuses to let us force that during functional construction
													A__   ,     A__		   =							self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				# Prepare our model
																				A__		   =							model_class(__UpperCAmelCase   )
																				# These are maximally general inputs for the model, with multiple None dimensions
																				# Hopefully this will catch any conditionals that fail for flexible shapes
																				A__		   =							{
																				    key: tf.keras.Input(shape=val.shape[1:]						,dtype=val.dtype						,name=__UpperCAmelCase   )
																				    for key, val in model.input_signature.items()
																				    if key in model.dummy_inputs
																				}
																				A__		   =							model(__UpperCAmelCase   )
																				self.assertTrue(outputs_dict is not None   )






def      UpperCAmelCase    (		):






							"""simple docstring"""





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


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





						@cached_property
						def  snake_case__		(				self   )    -> str:
													return (
													    EfficientFormerImageProcessor.from_pretrained('snap-research/efficientformer-l1-300'   )
													    if is_vision_available()
													    else None
													)





						@slow
						def  snake_case__		(				self   )    -> Optional[Any]:
													A__		   =							TFEfficientFormerForImageClassification.from_pretrained('snap-research/efficientformer-l1-300'   )
													A__		   =							self.default_image_processor
													A__		   =							prepare_img()
													A__		   =							image_processor(images=__UpperCAmelCase						,return_tensors='tf'   )
													# forward pass
													A__		   =							model(**__UpperCAmelCase						,training=__UpperCAmelCase   )
													# verify the logits
													A__		   =							tf.TensorShape((1, 10_00)   )
													self.assertEqual(outputs.logits.shape						,__UpperCAmelCase   )
													A__		   =							tf.constant([-0.0_5_5_5, 0.4_8_2_5, -0.0_8_5_2]   )
													self.assertTrue(np.allclose(outputs.logits[0, :3]						,__UpperCAmelCase						,atol=1e-4   )   )





						@slow
						def  snake_case__		(				self   )    -> Dict:
													A__		   =							TFEfficientFormerForImageClassificationWithTeacher.from_pretrained(
													    'snap-research/efficientformer-l1-300'   )
													A__		   =							self.default_image_processor
													A__		   =							prepare_img()
													A__		   =							image_processor(images=__UpperCAmelCase						,return_tensors='tf'   )
													# forward pass
													A__		   =							model(**__UpperCAmelCase						,training=__UpperCAmelCase   )
													# verify the logits
													A__		   =							tf.TensorShape((1, 10_00)   )
													self.assertEqual(outputs.logits.shape						,__UpperCAmelCase   )
													A__		   =							tf.constant([-0.1_3_1_2, 0.4_3_5_3, -1.0_4_9_9]   )
													self.assertTrue(np.allclose(outputs.logits[0, :3]						,__UpperCAmelCase						,atol=1e-4   )   )






"""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





