


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


lowerCAmelCase__				    =       logging.getLogger(__name__)
def  a__      (    SCREAMING_SNAKE_CASE   :   List[Any]=2    ,				SCREAMING_SNAKE_CASE   :   Tuple=3    ,				SCREAMING_SNAKE_CASE   :   Tuple=1_6    ,				SCREAMING_SNAKE_CASE   :   int = 1_0    ,				SCREAMING_SNAKE_CASE   :   int = 2    ):


  '''simple docstring'''







  def get_dataset(SCREAMING_SNAKE_CASE   :   Any    ):
    lowerCAmelCase					:					Tuple								=							torch.randn(batch_size * n_batches    ,				1    )
    return TensorDataset(SCREAMING_SNAKE_CASE    ,				a * x + b + 0.1 * torch.randn(batch_size * n_batches    ,				1    )    )

  lowerCAmelCase					:					List[Any]								=							get_dataset(SCREAMING_SNAKE_CASE    )
  lowerCAmelCase					:					str								=							get_dataset(SCREAMING_SNAKE_CASE    )
  lowerCAmelCase					:					Tuple								=							DataLoader(SCREAMING_SNAKE_CASE    ,				shuffle=SCREAMING_SNAKE_CASE    ,				batch_size=SCREAMING_SNAKE_CASE    ,				num_workers=4    )
  lowerCAmelCase					:					List[str]								=							DataLoader(SCREAMING_SNAKE_CASE    ,				shuffle=SCREAMING_SNAKE_CASE    ,				batch_size=SCREAMING_SNAKE_CASE    ,				num_workers=4    )
  return (train_dataloader, valid_dataloader)
def  a__      (    SCREAMING_SNAKE_CASE   :   str    ,				SCREAMING_SNAKE_CASE   :   Optional[Any]    ,				SCREAMING_SNAKE_CASE   :   Union[str, Any]    ,				SCREAMING_SNAKE_CASE   :   List[Any]    ,				SCREAMING_SNAKE_CASE   :   int    ,				SCREAMING_SNAKE_CASE   :   Dict=None    ):


  '''simple docstring'''







  lowerCAmelCase					:					int								=							[]
  for epoch in range(SCREAMING_SNAKE_CASE    ):
    # Train quickly
    model.train()
    for batch in dataloader:
      lowerCAmelCase						,  lowerCAmelCase					:					Any								=							batch
      lowerCAmelCase					:					Dict								=							model(SCREAMING_SNAKE_CASE    )
      lowerCAmelCase					:					int								=							torch.nn.functional.mse_loss(SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    )
      accelerator.backward(SCREAMING_SNAKE_CASE    )
      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    ):
       """simple docstring"""


       def __init__(						self						):






         """simple docstring"""





         super().__init__()
         lowerCAmelCase					:					Optional[int]								=							nn.Parameter(torch.randn(1						)						)
         lowerCAmelCase					:					Optional[int]								=							nn.Parameter(torch.randn(1						)						)







       def 	lowercase__							(						self		,					snake_case__						):






         """simple docstring"""





         return x * self.a + self.b


class 				SCREAMING_SNAKE_CASE__    (		unittest.TestCase    ):
       """simple docstring"""


       def 	lowercase__							(						self						):






         """simple docstring"""





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

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

       def 	lowercase__							(						self						):






         """simple docstring"""





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

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

           lowerCAmelCase					:					Dict								=							train(2		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__						)
           # Save everything
           lowerCAmelCase					:					Dict								=							os.path.join(snake_case__		,					"checkpoint"						)
           accelerator.save_state(snake_case__						)

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

       def 	lowercase__							(						self						):






         """simple docstring"""





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

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

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

           lowerCAmelCase					:					Union[str, Any]								=							train(2		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__						)
           # Save everything
           accelerator.save_state()

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

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Optional[Any]								=							torch.tensor([1, 2, 3]						)
         lowerCAmelCase					:					Tuple								=							torch.tensor([2, 3, 4]						)
         lowerCAmelCase					:					Tuple								=							DummyModel()
         lowerCAmelCase					:					Union[str, Any]								=							torch.optim.Adam(net.parameters()						)
         lowerCAmelCase					:					List[str]								=							Accelerator()
         with self.assertRaises(snake_case__						) as ve:
           accelerator.register_for_checkpointing(snake_case__		,					snake_case__		,					snake_case__		,					snake_case__						)
         lowerCAmelCase					:					List[str]								=							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 	lowercase__							(						self						):






         """simple docstring"""





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

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

       def 	lowercase__							(						self						):






         """simple docstring"""





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







       @require_cuda
       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					int								=							["torchrun", f"""--nproc_per_node={torch.cuda.device_count()}""", inspect.getfile(self.__class__						)]
         execute_subprocess_async(snake_case__		,					env=os.environ.copy()						)


if __name__ == "__main__":
     lowerCAmelCase__				    =       '''/tmp/accelerate/state_checkpointing'''
     lowerCAmelCase__				    =       DummyModel()
     lowerCAmelCase__				    =       torch.optim.Adam(params=model.parameters(), lr=1e-3)
     lowerCAmelCase__				    =       torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
     lowerCAmelCase__    ,lowerCAmelCase__				    =       dummy_dataloaders()
     lowerCAmelCase__				    =       ProjectConfiguration(automatic_checkpoint_naming=True)
     # Train baseline
     lowerCAmelCase__				    =       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)
     lowerCAmelCase__    ,lowerCAmelCase__    ,lowerCAmelCase__    ,lowerCAmelCase__    ,lowerCAmelCase__				    =       accelerator.prepare(
         model, optimizer, train_dataloader, valid_dataloader, scheduler
     )
     lowerCAmelCase__    ,lowerCAmelCase__				    =       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:
          lowerCAmelCase__				    =       group['''params'''][0].device
          break
     assert param_device.type == accelerator.device.type
     lowerCAmelCase__				    =       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:
          lowerCAmelCase__				    =       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:
          lowerCAmelCase__				    =       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__    :
       """simple docstring"""


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






         """simple docstring"""





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

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

         lowerCAmelCase					:					Tuple								=							None
         if self.use_labels:
           lowerCAmelCase					:					Any								=							ids_tensor([self.batch_size]		,					self.type_sequence_label_size						)

         lowerCAmelCase					:					Tuple								=							self.get_config()

         return config, pixel_values, labels

       def 	lowercase__							(						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=snake_case__		,					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 	lowercase__							(						self		,					snake_case__		,					snake_case__		,					snake_case__						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							TFEfficientFormerModel(config=snake_case__						)
         lowerCAmelCase					:					Optional[int]								=							model(snake_case__		,					training=snake_case__						)
         self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)						)

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__		,					snake_case__						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							self.type_sequence_label_size
         lowerCAmelCase					:					Dict								=							TFEfficientFormerForImageClassification(snake_case__						)
         lowerCAmelCase					:					Tuple								=							model(snake_case__		,					labels=snake_case__		,					training=snake_case__						)
         self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.type_sequence_label_size)						)

         # test greyscale images
         lowerCAmelCase					:					str								=							1
         lowerCAmelCase					:					Any								=							TFEfficientFormerForImageClassification(snake_case__						)

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







       def 	lowercase__							(						self						):






         """simple docstring"""





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


@require_tf
class 				SCREAMING_SNAKE_CASE__    (		lowercase       ,   lowercase       ,   unittest.TestCase    ):
       """simple docstring"""


       a  :    Optional[int]       =(
           (
               TFEfficientFormerModel,
               TFEfficientFormerForImageClassificationWithTeacher,
               TFEfficientFormerForImageClassification,
           )
           if is_tf_available()
           else ()
       )
       a  :    Union[str, Any]       =(
           {
               "feature-extraction": TFEfficientFormerModel,
               "image-classification": (
                   TFEfficientFormerForImageClassification,
                   TFEfficientFormerForImageClassificationWithTeacher,
               ),
           }
           if is_tf_available()
           else {}
       )

       a  :    int       =False

       a  :    Optional[Any]       =False
       a  :    List[Any]       =False
       a  :    str       =False
       a  :    List[Any]       =False

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Any								=							TFEfficientFormerModelTester(self						)
         lowerCAmelCase					:					Dict								=							ConfigTester(
             self		,					config_class=snake_case__		,					has_text_modality=snake_case__		,					hidden_size=37						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         self.config_tester.run_common_tests()

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






         """simple docstring"""





         pass

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






         """simple docstring"""





         pass

       def 	lowercase__							(						self						):






         """simple docstring"""





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

         for model_class in self.all_model_classes:
           lowerCAmelCase					:					Tuple								=							model_class(snake_case__						)
           lowerCAmelCase					:					str								=							inspect.signature(model.call						)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           lowerCAmelCase					:					Optional[int]								=							[*signature.parameters.keys()]

           lowerCAmelCase					:					str								=							["pixel_values"]
           self.assertListEqual(arg_names[:1]		,					snake_case__						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         def check_hidden_states_output(snake_case__		,					snake_case__		,					snake_case__						):
           lowerCAmelCase					:					List[Any]								=							model_class(snake_case__						)

           lowerCAmelCase					:					List[str]								=							model(**self._prepare_for_class(snake_case__		,					snake_case__						)		,					training=snake_case__						)
           lowerCAmelCase					:					Dict								=							outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

           lowerCAmelCase					:					List[Any]								=							getattr(
               self.model_tester		,					"expected_num_hidden_layers"		,					self.model_tester.num_hidden_layers + 1						)
           self.assertEqual(len(snake_case__						)		,					snake_case__						)

           if hasattr(self.model_tester		,					"encoder_seq_length"						):
             lowerCAmelCase					:					Union[str, Any]								=							self.model_tester.encoder_seq_length
             if hasattr(self.model_tester		,					"chunk_length"						) and self.model_tester.chunk_length > 1:
               lowerCAmelCase					:					Tuple								=							seq_length * self.model_tester.chunk_length
           else:
             lowerCAmelCase					:					List[str]								=							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:
             lowerCAmelCase					:					Tuple								=							outputs.decoder_hidden_states

             self.asseretIsInstance(snake_case__		,					(list, tuple)						)
             self.assertEqual(len(snake_case__						)		,					snake_case__						)
             lowerCAmelCase					:					int								=							getattr(self.model_tester		,					"seq_length"		,					snake_case__						)
             lowerCAmelCase					:					Tuple								=							getattr(self.model_tester		,					"decoder_seq_length"		,					snake_case__						)

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

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

         for model_class in self.all_model_classes:
           lowerCAmelCase					:					str								=							True
           check_hidden_states_output(snake_case__		,					snake_case__		,					snake_case__						)

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

           check_hidden_states_output(snake_case__		,					snake_case__		,					snake_case__						)

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__		,					snake_case__=False						):






         """simple docstring"""





         lowerCAmelCase					:					Optional[int]								=							super()._prepare_for_class(snake_case__		,					snake_case__		,					return_labels=snake_case__						)

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

         return inputs_dict

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Any								=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*snake_case__						)

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






         """simple docstring"""





         lowerCAmelCase					:					List[str]								=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*snake_case__						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					int								=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*snake_case__						)

       @slow
       def 	lowercase__							(						self						):






         """simple docstring"""





         for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowerCAmelCase					:					Optional[Any]								=							TFEfficientFormerModel.from_pretrained(snake_case__						)
           self.assertIsNotNone(snake_case__						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase					:					List[Any]								=							self.model_tester.prepare_config_and_inputs_for_common()
         lowerCAmelCase					:					int								=							True

         lowerCAmelCase					:					Union[str, Any]								=							getattr(self.model_tester		,					"seq_length"		,					snake_case__						)
         lowerCAmelCase					:					Dict								=							getattr(self.model_tester		,					"encoder_seq_length"		,					snake_case__						)
         lowerCAmelCase					:					Union[str, Any]								=							getattr(self.model_tester		,					"key_length"		,					snake_case__						)
         lowerCAmelCase					:					List[str]								=							getattr(self.model_tester		,					"chunk_length"		,					snake_case__						)

         if chunk_length is not None and hasattr(self.model_tester		,					"num_hashes"						):
           lowerCAmelCase					:					Dict								=							encoder_seq_length * self.model_tester.num_hashes

         for model_class in self.all_model_classes:
           lowerCAmelCase					:					int								=							True
           lowerCAmelCase					:					int								=							False
           lowerCAmelCase					:					Dict								=							True
           lowerCAmelCase					:					List[Any]								=							model_class(snake_case__						)

           lowerCAmelCase					:					Tuple								=							model(**self._prepare_for_class(snake_case__		,					snake_case__						)		,					training=snake_case__						)
           lowerCAmelCase					:					str								=							outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
           self.assertEqual(len(snake_case__						)		,					self.model_tester.num_attention_outputs						)

           # check that output_attentions also work using config
           del inputs_dict["output_attentions"]
           lowerCAmelCase					:					int								=							True
           lowerCAmelCase					:					Dict								=							model_class(snake_case__						)
           lowerCAmelCase					:					int								=							model(**self._prepare_for_class(snake_case__		,					snake_case__						)		,					training=snake_case__						)

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






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase					:					Tuple								=							self.model_tester.prepare_config_and_inputs_for_common()

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

def  a__      (    ):


  '''simple docstring'''







  lowerCAmelCase					:					Optional[Any]								=							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
  return image





@require_tf
@require_vision
class 				SCREAMING_SNAKE_CASE__    (		unittest.TestCase    ):
       """simple docstring"""


       @cached_property
       def 	lowercase__							(						self						):






         """simple docstring"""





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

       @slow
       def 	lowercase__							(						self						):






         """simple docstring"""





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







       @slow
       def 	lowercase__							(						self						):






         """simple docstring"""





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





