



"""simple docstring"""

import inspect
import logging
import os
import random
import shutil
import tempfile
import unittest

import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_cuda
from accelerate.utils import ProjectConfiguration, set_seed


lowercase__			   =  logging.getLogger(__name__)





def       __magic_name__    (   _lowerCamelCase				: Tuple=2  ,       _lowerCamelCase				: Union[str, Any]=3  ,       _lowerCamelCase				: Optional[int]=1_6  ,       _lowerCamelCase				: int = 1_0  ,       _lowerCamelCase				: int = 2					):

							def get_dataset(_lowerCamelCase				: str					):
														__a		:		Tuple									=       torch.randn(batch_size * n_batches  ,       1					)
														return TensorDataset(_lowerCamelCase  ,       a * x + b + 0.1 * torch.randn(batch_size * n_batches  ,       1					)					)

							__a		:		List[str]									=       get_dataset(_lowerCamelCase					)
							__a		:		int									=       get_dataset(_lowerCamelCase					)
							__a		:		Tuple									=       DataLoader(_lowerCamelCase  ,       shuffle=_lowerCamelCase  ,       batch_size=_lowerCamelCase  ,       num_workers=4					)
							__a		:		str									=       DataLoader(_lowerCamelCase  ,       shuffle=_lowerCamelCase  ,       batch_size=_lowerCamelCase  ,       num_workers=4					)
							return (train_dataloader, valid_dataloader)





def       __magic_name__    (   _lowerCamelCase				: Tuple  ,       _lowerCamelCase				: Optional[int]  ,       _lowerCamelCase				: str  ,       _lowerCamelCase				: Tuple  ,       _lowerCamelCase				: Union[str, Any]  ,       _lowerCamelCase				: Optional[Any]=None					):
							__a		:		Dict									=       []
							for epoch in range(_lowerCamelCase					):
														# Train quickly
														model.train()
														for batch in dataloader:
																					__a							,     __a		:		Tuple									=       batch
																					__a		:		Union[str, Any]									=       model(_lowerCamelCase					)
																					__a		:		int									=       torch.nn.functional.mse_loss(_lowerCamelCase  ,       _lowerCamelCase					)
																					accelerator.backward(_lowerCamelCase					)
																					optimizer.step()
																					optimizer.zero_grad()
														rands.append(random.random()					)  # Introduce some randomness
														if scheduler is not None:
																					scheduler.step()
							return rands




class SCREAMING_SNAKE_CASE__					(		nn.Module      ):



						def __init__(self							):





													'''simple docstring'''

													super().__init__()
													__a		:		int									=       nn.Parameter(torch.randn(1							)							)
													__a		:		Any									=       nn.Parameter(torch.randn(1							)							)





						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

													return x * self.a + self.b







class SCREAMING_SNAKE_CASE__					(		unittest.TestCase      ):



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													with tempfile.TemporaryDirectory() as tmpdir:
																				set_seed(42							)
																				__a		:		Optional[int]									=       DummyModel()
																				__a		:		List[str]									=       torch.optim.Adam(params=model.parameters()						,     lr=1e-3							)
																				__a							,     __a		:		Optional[int]									=       dummy_dataloaders()
																				__a		:		List[str]									=       ProjectConfiguration(total_limit=1						,     project_dir=_lowercase						,     automatic_checkpoint_naming=_lowercase							)
																				# Train baseline
																				__a		:		str									=       Accelerator(project_config=_lowercase							)
																				__a							,     __a							,     __a							,     __a		:		Optional[int]									=       accelerator.prepare(
																				    _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				# Save initial
																				accelerator.save_state()

																				# Save second state
																				accelerator.save_state()
																				self.assertEqual(len(os.listdir(accelerator.project_dir							)							)						,     1							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													with tempfile.TemporaryDirectory() as tmpdir:
																				set_seed(42							)
																				__a		:		Optional[int]									=       DummyModel()
																				__a		:		List[str]									=       torch.optim.Adam(params=model.parameters()						,     lr=1e-3							)
																				__a							,     __a		:		Dict									=       dummy_dataloaders()
																				# Train baseline
																				__a		:		Dict									=       Accelerator()
																				__a							,     __a							,     __a							,     __a		:		Dict									=       accelerator.prepare(
																				    _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				# Save initial
																				__a		:		Tuple									=       os.path.join(_lowercase						,     """initial"""							)
																				accelerator.save_state(_lowercase							)
																				((__a)							,     (__a))		:		Tuple									=       model.a.item(), model.b.item()
																				__a		:		Optional[int]									=       optimizer.state_dict()
																				__a		:		Tuple									=       train(3						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				((__a)							,     (__a))		:		List[Any]									=       model.a.item(), model.b.item()
																				__a		:		List[str]									=       optimizer.state_dict()

																				# Train partially
																				set_seed(42							)
																				__a		:		Tuple									=       DummyModel()
																				__a		:		Any									=       torch.optim.Adam(params=model.parameters()						,     lr=1e-3							)
																				__a							,     __a		:		int									=       dummy_dataloaders()
																				__a		:		Optional[Any]									=       Accelerator()
																				__a							,     __a							,     __a							,     __a		:		List[Any]									=       accelerator.prepare(
																				    _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				accelerator.load_state(_lowercase							)
																				((__a)							,     (__a))		:		Union[str, Any]									=       model.a.item(), model.b.item()
																				__a		:		List[Any]									=       optimizer.state_dict()
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)

																				__a		:		Optional[Any]									=       train(2						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				# Save everything
																				__a		:		Optional[Any]									=       os.path.join(_lowercase						,     """checkpoint"""							)
																				accelerator.save_state(_lowercase							)

																				# Load everything back in and make sure all states work
																				accelerator.load_state(_lowercase							)
																				test_rands += train(1						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				((__a)							,     (__a))		:		str									=       model.a.item(), model.b.item()
																				__a		:		int									=       optimizer.state_dict()
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													with tempfile.TemporaryDirectory() as tmpdir:
																				set_seed(42							)
																				__a		:		List[Any]									=       DummyModel()
																				__a		:		Any									=       torch.optim.Adam(params=model.parameters()						,     lr=1e-3							)
																				__a							,     __a		:		List[str]									=       dummy_dataloaders()
																				__a		:		Optional[Any]									=       ProjectConfiguration(automatic_checkpoint_naming=_lowercase							)

																				# Train baseline
																				__a		:		List[str]									=       Accelerator(project_dir=_lowercase						,     project_config=_lowercase							)
																				__a							,     __a							,     __a							,     __a		:		Union[str, Any]									=       accelerator.prepare(
																				    _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				# Save initial
																				accelerator.save_state()
																				((__a)							,     (__a))		:		Dict									=       model.a.item(), model.b.item()
																				__a		:		Optional[Any]									=       optimizer.state_dict()
																				__a		:		str									=       train(3						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				((__a)							,     (__a))		:		Union[str, Any]									=       model.a.item(), model.b.item()
																				__a		:		int									=       optimizer.state_dict()

																				# Train partially
																				set_seed(42							)
																				__a		:		Union[str, Any]									=       DummyModel()
																				__a		:		List[Any]									=       torch.optim.Adam(params=model.parameters()						,     lr=1e-3							)
																				__a							,     __a		:		List[Any]									=       dummy_dataloaders()
																				__a		:		int									=       ProjectConfiguration(iteration=1						,     automatic_checkpoint_naming=_lowercase							)
																				__a		:		str									=       Accelerator(project_dir=_lowercase						,     project_config=_lowercase							)
																				__a							,     __a							,     __a							,     __a		:		Dict									=       accelerator.prepare(
																				    _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				accelerator.load_state(os.path.join(_lowercase						,     """checkpoints"""						,     """checkpoint_0"""							)							)
																				((__a)							,     (__a))		:		Optional[Any]									=       model.a.item(), model.b.item()
																				__a		:		Union[str, Any]									=       optimizer.state_dict()
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)

																				__a		:		Dict									=       train(2						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				# Save everything
																				accelerator.save_state()

																				# Load everything back in and make sure all states work
																				accelerator.load_state(os.path.join(_lowercase						,     """checkpoints"""						,     """checkpoint_1"""							)							)
																				test_rands += train(1						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				((__a)							,     (__a))		:		int									=       model.a.item(), model.b.item()
																				__a		:		str									=       optimizer.state_dict()
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)
																				self.assertEqual(_lowercase						,     _lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		str									=       torch.tensor([1, 2, 3]							)
													__a		:		Any									=       torch.tensor([2, 3, 4]							)
													__a		:		List[str]									=       DummyModel()
													__a		:		Optional[Any]									=       torch.optim.Adam(net.parameters()							)
													__a		:		List[str]									=       Accelerator()
													with self.assertRaises(_lowercase							) as ve:
																				accelerator.register_for_checkpointing(_lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
													__a		:		List[Any]									=       str(ve.exception							)
													self.assertTrue("""Item at index 0""" in message							)
													self.assertTrue("""Item at index 1""" in message							)
													self.assertFalse("""Item at index 2""" in message							)
													self.assertFalse("""Item at index 3""" in message							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													with tempfile.TemporaryDirectory() as tmpdir:
																				set_seed(42							)
																				__a		:		int									=       DummyModel()
																				__a		:		int									=       torch.optim.Adam(params=model.parameters()						,     lr=1e-3							)
																				__a		:		Tuple									=       torch.optim.lr_scheduler.StepLR(_lowercase						,     step_size=1						,     gamma=0.99							)
																				__a							,     __a		:		Optional[Any]									=       dummy_dataloaders()
																				__a		:		Optional[int]									=       ProjectConfiguration(automatic_checkpoint_naming=_lowercase							)
																				# Train baseline
																				__a		:		List[Any]									=       Accelerator(project_dir=_lowercase						,     project_config=_lowercase							)
																				__a							,     __a							,     __a							,     __a							,     __a		:		Union[str, Any]									=       accelerator.prepare(
																				    _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				# Save initial
																				accelerator.save_state()
																				__a		:		Optional[int]									=       scheduler.state_dict()
																				train(3						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							)
																				self.assertNotEqual(_lowercase						,     scheduler.state_dict()							)

																				# Load everything back in and make sure all states work
																				accelerator.load_state(os.path.join(_lowercase						,     """checkpoints"""						,     """checkpoint_0"""							)							)
																				self.assertEqual(_lowercase						,     scheduler.state_dict()							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													with tempfile.TemporaryDirectory() as tmpdir:
																				set_seed(42							)
																				__a		:		Dict									=       DummyModel()
																				__a		:		Dict									=       ProjectConfiguration(automatic_checkpoint_naming=_lowercase						,     total_limit=2							)
																				# Train baseline
																				__a		:		Tuple									=       Accelerator(project_dir=_lowercase						,     project_config=_lowercase							)
																				__a		:		str									=       accelerator.prepare(_lowercase							)
																				# Save 3 states:
																				for _ in range(11							):
																											accelerator.save_state()
																				self.assertTrue(not os.path.exists(os.path.join(_lowercase						,     """checkpoints"""						,     """checkpoint_0"""							)							)							)
																				self.assertTrue(os.path.exists(os.path.join(_lowercase						,     """checkpoints"""						,     """checkpoint_9"""							)							)							)
																				self.assertTrue(os.path.exists(os.path.join(_lowercase						,     """checkpoints"""						,     """checkpoint_10"""							)							)							)





						@require_cuda
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		List[Any]									=       ["""torchrun""", F'''--nproc_per_node={torch.cuda.device_count()}''', inspect.getfile(self.__class__							)]
													execute_subprocess_async(_lowercase						,     env=os.environ.copy()							)


if __name__ == "__main__":
	lowercase__			   =  "/tmp/accelerate/state_checkpointing"
	lowercase__			   =  DummyModel()
	lowercase__			   =  torch.optim.Adam(params=model.parameters(), lr=1e-3)
	lowercase__			   =  torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
	lowercase__   ,   lowercase__			   =  dummy_dataloaders()
	lowercase__			   =  ProjectConfiguration(automatic_checkpoint_naming=True)
	# Train baseline
	lowercase__			   =  Accelerator(project_dir=savedir, project_config=project_config, mixed_precision="no")
	if accelerator.process_index == 0:
		if os.path.exists(savedir):
			shutil.rmtree(savedir)
		os.makedirs(savedir)
	lowercase__   ,   lowercase__   ,   lowercase__   ,   lowercase__   ,   lowercase__			   =  accelerator.prepare(
	    model, optimizer, train_dataloader, valid_dataloader, scheduler
	)
	lowercase__   ,   lowercase__			   =  accelerator.prepare(model, optimizer)
	train(3, model, train_dataloader, optimizer, accelerator, scheduler)
	# Check that the intial optimizer is loaded on the GPU
	for group in optimizer.param_groups:
		lowercase__			   =  group["params"][0].device
		break
	assert param_device.type == accelerator.device.type
	lowercase__			   =  model.cpu()
	accelerator.wait_for_everyone()
	accelerator.save_state()
	accelerator.wait_for_everyone()

	# Check CPU state
	accelerator.load_state(os.path.join(savedir, "checkpoints", "checkpoint_0"), map_location="cpu")
	for group in optimizer.param_groups:
		lowercase__			   =  group["params"][0].device
		break
	assert (
	    param_device.type == torch.device("cpu").type
	), f"Loaded optimizer states did not match, expected to be loaded on the CPU but got {param_device}"

	# Check device state
	model.to(accelerator.device)
	accelerator.load_state(os.path.join(savedir, "checkpoints", "checkpoint_0"), map_location="on_device")
	for group in optimizer.param_groups:
		lowercase__			   =  group["params"][0].device
		break
	assert (
	    param_device.type == accelerator.device.type
	), f"Loaded optimizer states did not match, expected to be loaded on {accelerator.device} but got {param_device}"

	# Check error
	with pytest.raises(TypeError, match="Unsupported optimizer map location passed"):
		accelerator.load_state(os.path.join(savedir, "checkpoints", "checkpoint_0"), map_location="invalid")
	accelerator.wait_for_everyone()
	if accelerator.process_index == 0:
		shutil.rmtree(savedir)
	accelerator.wait_for_everyone()






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



						def __init__(self						,     _lowercase						,     _lowercase = 13						,     _lowercase = 64						,     _lowercase = 2						,     _lowercase = 3						,     _lowercase = 3						,     _lowercase = True						,     _lowercase = True						,     _lowercase = 128						,     _lowercase=[16, 32, 64, 128]						,     _lowercase = 7						,     _lowercase = 4						,     _lowercase = 37						,     _lowercase = "gelu"						,     _lowercase = 0.1						,     _lowercase = 0.1						,     _lowercase = 10						,     _lowercase = 0.02						,     _lowercase = 2						,     _lowercase = 1						,     _lowercase = 128						,     _lowercase = [2, 2, 2, 2]						,     _lowercase = 2						,     _lowercase = 2						,     ):





													'''simple docstring'''

													__a		:		str									=       parent
													__a		:		List[Any]									=       batch_size
													__a		:		int									=       image_size
													__a		:		Tuple									=       patch_size
													__a		:		str									=       num_channels
													__a		:		Union[str, Any]									=       is_training
													__a		:		List[Any]									=       use_labels
													__a		:		int									=       hidden_size
													__a		:		Optional[Any]									=       num_hidden_layers
													__a		:		List[Any]									=       num_attention_heads
													__a		:		Dict									=       intermediate_size
													__a		:		str									=       hidden_act
													__a		:		Dict									=       hidden_dropout_prob
													__a		:		str									=       attention_probs_dropout_prob
													__a		:		Optional[int]									=       type_sequence_label_size
													__a		:		Dict									=       initializer_range
													__a		:		Dict									=       encoder_stride
													__a		:		int									=       num_attention_outputs
													__a		:		List[Any]									=       embed_dim
													__a		:		Optional[Any]									=       embed_dim + 1
													__a		:		Optional[Any]									=       resolution
													__a		:		Optional[Any]									=       depths
													__a		:		Union[str, Any]									=       hidden_sizes
													__a		:		List[str]									=       dim
													__a		:		Any									=       mlp_expansion_ratio



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Union[str, Any]									=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

													__a		:		str									=       None
													if self.use_labels:
																				__a		:		List[str]									=       ids_tensor([self.batch_size]						,     self.type_sequence_label_size							)

													__a		:		List[str]									=       self.get_config()

													return config, pixel_values, labels



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													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=_lowercase						,     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     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		Optional[Any]									=       TFEfficientFormerModel(config=_lowercase							)
													__a		:		List[Any]									=       model(_lowercase						,     training=_lowercase							)
													self.parent.assertEqual(result.last_hidden_state.shape						,     (self.batch_size, self.seq_length, self.hidden_size)							)



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		Optional[Any]									=       self.type_sequence_label_size
													__a		:		Any									=       TFEfficientFormerForImageClassification(_lowercase							)
													__a		:		Union[str, Any]									=       model(_lowercase						,     labels=_lowercase						,     training=_lowercase							)
													self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.type_sequence_label_size)							)

													# test greyscale images
													__a		:		Optional[Any]									=       1
													__a		:		int									=       TFEfficientFormerForImageClassification(_lowercase							)

													__a		:		List[str]									=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
													__a		:		str									=       model(_lowercase						,     labels=_lowercase							)
													self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.type_sequence_label_size)							)





						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Any									=       self.prepare_config_and_inputs()
													__a							,     __a							,     __a		:		Tuple									=       config_and_inputs
													__a		:		Tuple									=       {"""pixel_values""": pixel_values}
													return config, inputs_dict







@require_tf
class SCREAMING_SNAKE_CASE__					(		__snake_case     ,					__snake_case     ,					unittest.TestCase      ):
						_lowerCAmelCase							=					(
						    (
						        TFEfficientFormerModel,
						        TFEfficientFormerForImageClassificationWithTeacher,
						        TFEfficientFormerForImageClassification,
						    )
						    if is_tf_available()
						    else ()
						)
						_lowerCAmelCase							=					(
						    {
						        "feature-extraction": TFEfficientFormerModel,
						        "image-classification": (
						            TFEfficientFormerForImageClassification,
						            TFEfficientFormerForImageClassificationWithTeacher,
						        ),
						    }
						    if is_tf_available()
						    else {}
						)

						_lowerCAmelCase							=					False

						_lowerCAmelCase							=					False
						_lowerCAmelCase							=					False
						_lowerCAmelCase							=					False
						_lowerCAmelCase							=					False



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Tuple									=       TFEfficientFormerModelTester(self							)
													__a		:		Any									=       ConfigTester(
													    self						,     config_class=_lowercase						,     has_text_modality=_lowercase						,     hidden_size=37							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													self.config_tester.run_common_tests()



						@unittest.skip(reason="""EfficientFormer does not use inputs_embeds"""							)
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													pass



						@unittest.skip(reason="""EfficientFormer does not support input and output embeddings"""							)
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													pass



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

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

													for model_class in self.all_model_classes:
																				__a		:		Dict									=       model_class(_lowercase							)
																				__a		:		Optional[Any]									=       inspect.signature(model.call							)
																				# signature.parameters is an OrderedDict => so arg_names order is deterministic
																				__a		:		Optional[Any]									=       [*signature.parameters.keys()]

																				__a		:		Union[str, Any]									=       ["""pixel_values"""]
																				self.assertListEqual(arg_names[:1]						,     _lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													def check_hidden_states_output(_lowercase						,     _lowercase						,     _lowercase							):
																				__a		:		Tuple									=       model_class(_lowercase							)

																				__a		:		int									=       model(**self._prepare_for_class(_lowercase						,     _lowercase							)						,     training=_lowercase							)
																				__a		:		Tuple									=       outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																				__a		:		str									=       getattr(
																				    self.model_tester						,     """expected_num_hidden_layers"""						,     self.model_tester.num_hidden_layers + 1							)
																				self.assertEqual(len(_lowercase							)						,     _lowercase							)

																				if hasattr(self.model_tester						,     """encoder_seq_length"""							):
																											__a		:		Any									=       self.model_tester.encoder_seq_length
																											if hasattr(self.model_tester						,     """chunk_length"""							) and self.model_tester.chunk_length > 1:
																																		__a		:		int									=       seq_length * self.model_tester.chunk_length
																				else:
																											__a		:		Any									=       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		:		Optional[int]									=       outputs.decoder_hidden_states

																											self.asseretIsInstance(_lowercase						,     (list, tuple)							)
																											self.assertEqual(len(_lowercase							)						,     _lowercase							)
																											__a		:		Any									=       getattr(self.model_tester						,     """seq_length"""						,     _lowercase							)
																											__a		:		List[Any]									=       getattr(self.model_tester						,     """decoder_seq_length"""						,     _lowercase							)

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

													__a							,     __a		:		Union[str, Any]									=       self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				__a		:		Dict									=       True
																				check_hidden_states_output(_lowercase						,     _lowercase						,     _lowercase							)

																				# check that output_hidden_states also work using config
																				del inputs_dict["output_hidden_states"]
																				__a		:		int									=       True

																				check_hidden_states_output(_lowercase						,     _lowercase						,     _lowercase							)



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase=False							):





													'''simple docstring'''

													__a		:		Any									=       super()._prepare_for_class(_lowercase						,     _lowercase						,     return_labels=_lowercase							)

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

													return inputs_dict



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Optional[Any]									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*_lowercase							)



						@unittest.skip(reason="""EfficientFormer does not implement masked image modeling yet"""							)
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Union[str, Any]									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_image_modeling(*_lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Optional[Any]									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_image_classification(*_lowercase							)



						@slow
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				__a		:		Union[str, Any]									=       TFEfficientFormerModel.from_pretrained(_lowercase							)
																				self.assertIsNotNone(_lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a							,     __a		:		Optional[Any]									=       self.model_tester.prepare_config_and_inputs_for_common()
													__a		:		int									=       True

													__a		:		Optional[int]									=       getattr(self.model_tester						,     """seq_length"""						,     _lowercase							)
													__a		:		Dict									=       getattr(self.model_tester						,     """encoder_seq_length"""						,     _lowercase							)
													__a		:		Dict									=       getattr(self.model_tester						,     """key_length"""						,     _lowercase							)
													__a		:		int									=       getattr(self.model_tester						,     """chunk_length"""						,     _lowercase							)

													if chunk_length is not None and hasattr(self.model_tester						,     """num_hashes"""							):
																				__a		:		List[str]									=       encoder_seq_length * self.model_tester.num_hashes

													for model_class in self.all_model_classes:
																				__a		:		List[Any]									=       True
																				__a		:		Tuple									=       False
																				__a		:		List[Any]									=       True
																				__a		:		int									=       model_class(_lowercase							)

																				__a		:		List[Any]									=       model(**self._prepare_for_class(_lowercase						,     _lowercase							)						,     training=_lowercase							)
																				__a		:		Dict									=       outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
																				self.assertEqual(len(_lowercase							)						,     self.model_tester.num_attention_outputs							)

																				# check that output_attentions also work using config
																				del inputs_dict["output_attentions"]
																				__a		:		Optional[Any]									=       True
																				__a		:		List[str]									=       model_class(_lowercase							)
																				__a		:		Dict									=       model(**self._prepare_for_class(_lowercase						,     _lowercase							)						,     training=_lowercase							)

																				__a		:		int									=       outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
																				self.assertEqual(len(_lowercase							)						,     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     lowerCAmelCase__(self							):





													'''simple docstring'''

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

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







def       __magic_name__    (   ):
							__a		:		int									=       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
							return image







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



						@cached_property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return (
													    EfficientFormerImageProcessor.from_pretrained("""snap-research/efficientformer-l1-300"""							)
													    if is_vision_available()
													    else None
													)



						@slow
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		str									=       TFEfficientFormerForImageClassification.from_pretrained("""snap-research/efficientformer-l1-300"""							)
													__a		:		Optional[Any]									=       self.default_image_processor
													__a		:		List[str]									=       prepare_img()
													__a		:		int									=       image_processor(images=_lowercase						,     return_tensors="""tf"""							)
													# forward pass
													__a		:		Optional[Any]									=       model(**_lowercase						,     training=_lowercase							)
													# verify the logits
													__a		:		str									=       tf.TensorShape((1, 1000)							)
													self.assertEqual(outputs.logits.shape						,     _lowercase							)
													__a		:		Dict									=       tf.constant([-0.0555, 0.4825, -0.0852]							)
													self.assertTrue(np.allclose(outputs.logits[0, :3]						,     _lowercase						,     atol=1e-4							)							)





						@slow
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Any									=       TFEfficientFormerForImageClassificationWithTeacher.from_pretrained(
													    """snap-research/efficientformer-l1-300"""							)
													__a		:		Any									=       self.default_image_processor
													__a		:		str									=       prepare_img()
													__a		:		str									=       image_processor(images=_lowercase						,     return_tensors="""tf"""							)
													# forward pass
													__a		:		List[Any]									=       model(**_lowercase						,     training=_lowercase							)
													# verify the logits
													__a		:		int									=       tf.TensorShape((1, 1000)							)
													self.assertEqual(outputs.logits.shape						,     _lowercase							)
													__a		:		List[str]									=       tf.constant([-0.1312, 0.4353, -1.0499]							)
													self.assertTrue(np.allclose(outputs.logits[0, :3]						,     _lowercase						,     atol=1e-4							)							)

