





"""simple docstring"""


import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin





def 							a__			(  _UpperCAmelCase=2	,     _UpperCAmelCase=True	,     _UpperCAmelCase=False	)    ->  Optional[Any]:

 '''simple docstring'''






 snake_case__    							=				Accelerator(step_scheduler_with_optimizer=_UpperCAmelCase	,     split_batches=_UpperCAmelCase	)
 snake_case__    							=				torch.nn.Linear(2	,     4	)
 snake_case__    							=				torch.optim.AdamW(model.parameters()	,     lr=1.0	)
 snake_case__    							=				torch.optim.lr_scheduler.OneCycleLR(_UpperCAmelCase	,     max_lr=0.01	,     steps_per_epoch=2	,     epochs=1	)
 snake_case__							,   snake_case__							,   snake_case__    							=				accelerator.prepare(_UpperCAmelCase	,     _UpperCAmelCase	,     _UpperCAmelCase	)

 # Optimizer has stepped
 scheduler.step()
 if step_scheduler_with_optimizer or (num_processes == 1):
  assert (
      scheduler.scheduler.last_epoch == num_processes
  ), f"Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})"
 else:
  assert (
      scheduler.scheduler.last_epoch != num_processes
  ), f"Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})"





def 							a__			(  _UpperCAmelCase=2	,     _UpperCAmelCase=True	,     _UpperCAmelCase=False	)    ->  Union[str, Any]:

 '''simple docstring'''






 snake_case__    							=				Accelerator(step_scheduler_with_optimizer=_UpperCAmelCase	,     split_batches=_UpperCAmelCase	)
 snake_case__    							=				torch.nn.Linear(2	,     4	)
 snake_case__    							=				torch.optim.AdamW(model.parameters()	,     lr=1.0	)
 snake_case__    							=				torch.optim.lr_scheduler.LambdaLR(_UpperCAmelCase	,     lr_lambda=lambda _UpperCAmelCase	: 1 - n / 1_0	)
 snake_case__							,   snake_case__							,   snake_case__    							=				accelerator.prepare(_UpperCAmelCase	,     _UpperCAmelCase	,     _UpperCAmelCase	)

 # Optimizer has stepped
 snake_case__    							=				False
 scheduler.step()
 snake_case__    							=				1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 1_0
 assert (
     scheduler.get_last_lr()[0] == expected_lr
 ), f"Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"

 # Optimizer has not stepped
 snake_case__    							=				True
 scheduler.step()
 if not step_scheduler_with_optimizer:
  snake_case__    							=				1 - 2 / 1_0
 assert (
     scheduler.get_last_lr()[0] == expected_lr
 ), f"Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"





def 							a__			(  _UpperCAmelCase = 2	)    ->  Tuple:

 '''simple docstring'''






 from transformers import get_linear_schedule_with_warmup

 snake_case__    							=				[1, 2, 4]
 for num_steps in steps:
  snake_case__    							=				GradientAccumulationPlugin(num_steps=_UpperCAmelCase	,     adjust_scheduler=num_steps > 1	)
  snake_case__    							=				Accelerator(gradient_accumulation_plugin=_UpperCAmelCase	)
  snake_case__    							=				torch.nn.Linear(2	,     4	)
  snake_case__    							=				torch.optim.AdamW(model.parameters()	,     lr=10.0	)
  snake_case__    							=				get_linear_schedule_with_warmup(optimizer=_UpperCAmelCase	,     num_warmup_steps=0	,     num_training_steps=2_0	)

  snake_case__							,   snake_case__							,   snake_case__    							=				accelerator.prepare(_UpperCAmelCase	,     _UpperCAmelCase	,     _UpperCAmelCase	)

  for i in range(1_0 * num_steps	):
   with accelerator.accumulate(_UpperCAmelCase	):
    optimizer.step()
    scheduler.step()

   if i == (1_0 * num_steps - 2):
    assert (
        scheduler.get_last_lr()[0] != 0
    ), f"Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}"
  assert (
      scheduler.get_last_lr()[0] == 0
  ), f"Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}"
  GradientState._reset_state()



@require_cpu
class      lowercase_				(     unittest.TestCase							):

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1     )   ,					num_processes=1     )
        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1   ,					split_batches=__SCREAMING_SNAKE_CASE     )   ,					num_processes=1     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1     )   ,					num_processes=1     )
        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1   ,					split_batches=__SCREAMING_SNAKE_CASE     )   ,					num_processes=1     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1   ,					step_scheduler_with_optimizer=__SCREAMING_SNAKE_CASE     )   ,					num_processes=1     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1   ,					step_scheduler_with_optimizer=__SCREAMING_SNAKE_CASE     )   ,					num_processes=1     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE     )
        debug_launcher(__SCREAMING_SNAKE_CASE     )
        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1   ,					split_batches=__SCREAMING_SNAKE_CASE     )   ,					num_processes=1     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE     )
        debug_launcher(__SCREAMING_SNAKE_CASE     )
        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1   ,					split_batches=__SCREAMING_SNAKE_CASE     )   ,					num_processes=1     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE     )
        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					step_scheduler_with_optimizer=__SCREAMING_SNAKE_CASE     )     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE     )
        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					step_scheduler_with_optimizer=__SCREAMING_SNAKE_CASE     )     )

       @require_huggingface_suite
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        AcceleratorState._reset_state(__SCREAMING_SNAKE_CASE     )
        debug_launcher(partial(__SCREAMING_SNAKE_CASE   ,					num_processes=1     )     )
        debug_launcher(__SCREAMING_SNAKE_CASE     )








"""simple docstring"""


import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import MODEL_MAPPING, PoolFormerConfig, PoolFormerForImageClassification, PoolFormerModel
   from transformers.models.poolformer.modeling_poolformer import POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import PoolFormerImageProcessor



class      lowercase_				(     __lowerCAmelCase							):

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.config_class(**self.inputs_dict     )
        self.parent.assertTrue(hasattr(__SCREAMING_SNAKE_CASE   ,					"hidden_sizes"     )     )
        self.parent.assertTrue(hasattr(__SCREAMING_SNAKE_CASE   ,					"num_encoder_blocks"     )     )



class      lowercase_				:

       def __init__(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE=13   ,					__SCREAMING_SNAKE_CASE=64   ,					__SCREAMING_SNAKE_CASE=3   ,					__SCREAMING_SNAKE_CASE=4   ,					__SCREAMING_SNAKE_CASE=[2, 2, 2, 2]   ,					__SCREAMING_SNAKE_CASE=[8, 4, 2, 1]   ,					__SCREAMING_SNAKE_CASE=[16, 32, 64, 128]   ,					__SCREAMING_SNAKE_CASE=[1, 4, 8, 16]   ,					__SCREAMING_SNAKE_CASE=False   ,					__SCREAMING_SNAKE_CASE=True   ,					__SCREAMING_SNAKE_CASE="gelu"   ,					__SCREAMING_SNAKE_CASE=0.1   ,					__SCREAMING_SNAKE_CASE=0.02   ,					__SCREAMING_SNAKE_CASE=3   ,					__SCREAMING_SNAKE_CASE=None   ,					):






        '''simple docstring'''

        snake_case__    							=				parent
        snake_case__    							=				batch_size
        snake_case__    							=				image_size
        snake_case__    							=				num_channels
        snake_case__    							=				num_encoder_blocks
        snake_case__    							=				sr_ratios
        snake_case__    							=				depths
        snake_case__    							=				hidden_sizes
        snake_case__    							=				downsampling_rates
        snake_case__    							=				is_training
        snake_case__    							=				use_labels
        snake_case__    							=				hidden_act
        snake_case__    							=				hidden_dropout_prob
        snake_case__    							=				initializer_range
        snake_case__    							=				num_labels
        snake_case__    							=				scope

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

        snake_case__    							=				None
        if self.use_labels:
         snake_case__    							=				ids_tensor([self.batch_size, self.image_size, self.image_size]   ,					self.num_labels     )

        snake_case__    							=				PoolFormerConfig(
            image_size=self.image_size   ,					num_channels=self.num_channels   ,					num_encoder_blocks=self.num_encoder_blocks   ,					depths=self.depths   ,					hidden_sizes=self.hidden_sizes   ,					hidden_act=self.hidden_act   ,					hidden_dropout_prob=self.hidden_dropout_prob   ,					initializer_range=self.initializer_range   ,					)

        return config, pixel_values, labels

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        snake_case__    							=				PoolFormerModel(config=__SCREAMING_SNAKE_CASE     )
        model.to(__SCREAMING_SNAKE_CASE     )
        model.eval()
        snake_case__    							=				model(__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				snake_case__    							=				self.image_size // 32.0
        self.parent.assertEqual(
            result.last_hidden_state.shape   ,					(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

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



@require_torch
class      lowercase_				(     __lowerCAmelCase				,__lowerCAmelCase				,unittest.TestCase							):
       snake_case__				:							Optional[Any]											=     (PoolFormerModel, PoolFormerForImageClassification) if is_torch_available() else ()
       snake_case__				:							Dict											=     (
           {"""feature-extraction""": PoolFormerModel, """image-classification""": PoolFormerForImageClassification}
           if is_torch_available()
           else {}
       )

       snake_case__				:							Optional[int]											=     False
       snake_case__				:							Optional[Any]											=     False
       snake_case__				:							List[Any]											=     False
       snake_case__				:							Any											=     False
       snake_case__				:							Tuple											=     False

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				PoolFormerModelTester(self     )
        snake_case__    							=				PoolFormerConfigTester(self   ,					config_class=__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        self.config_tester.run_common_tests()

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE     )

       @unittest.skip("PoolFormer does not use inputs_embeds"     )
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        pass

       @unittest.skip("PoolFormer does not have get_input_embeddings method and get_output_embeddings methods"     )
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        pass

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        def check_hidden_states_output(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ):
         snake_case__    							=				model_class(__SCREAMING_SNAKE_CASE     )
         model.to(__SCREAMING_SNAKE_CASE     )
         model.eval()

         with torch.no_grad():
          snake_case__    							=				model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     )     )

         snake_case__    							=				outputs.hidden_states

         snake_case__    							=				self.model_tester.num_encoder_blocks
         self.assertEqual(len(__SCREAMING_SNAKE_CASE     )   ,					__SCREAMING_SNAKE_CASE     )

         # verify the first hidden states (first block)
         self.assertListEqual(
             list(hidden_states[0].shape[-3:]     )   ,					[
                 self.model_tester.hidden_sizes[0],
                 self.model_tester.image_size // 4,
                 self.model_tester.image_size // 4,
             ]   ,					)

        snake_case__							,   snake_case__    							=				self.model_tester.prepare_config_and_inputs_for_common()

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

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

         check_hidden_states_output(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        if not self.model_tester.is_training:
         return

        snake_case__							,   snake_case__    							=				self.model_tester.prepare_config_and_inputs_for_common()
        snake_case__    							=				True

        for model_class in self.all_model_classes:
         if model_class in get_values(__SCREAMING_SNAKE_CASE     ):
          continue
         snake_case__    							=				model_class(__SCREAMING_SNAKE_CASE     )
         model.to(__SCREAMING_SNAKE_CASE     )
         model.train()
         snake_case__    							=				self._prepare_for_class(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					return_labels=__SCREAMING_SNAKE_CASE     )
         snake_case__    							=				model(**__SCREAMING_SNAKE_CASE     ).loss
         loss.backward()

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__							,   snake_case__    							=				self.model_tester.prepare_config_and_inputs_for_common()

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

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

       @slow
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        for model_name in POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case__    							=				PoolFormerModel.from_pretrained(__SCREAMING_SNAKE_CASE     )
         self.assertIsNotNone(__SCREAMING_SNAKE_CASE     )





def 							a__			(  )    ->  int:

 '''simple docstring'''






 snake_case__    							=				Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
 return image



@require_torch
class      lowercase_				(     unittest.TestCase							):

       @slow
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        snake_case__    							=				PoolFormerImageProcessor()
        snake_case__    							=				PoolFormerForImageClassification.from_pretrained("sail/poolformer_s12"     ).to(__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				image_processor(images=prepare_img()   ,					return_tensors="pt"     ).to(__SCREAMING_SNAKE_CASE     )

        # forward pass
        with torch.no_grad():
         snake_case__    							=				model(**__SCREAMING_SNAKE_CASE     )

        # verify the logits
        snake_case__    							=				torch.Size((1, 1000)     )
        self.assertEqual(outputs.logits.shape   ,					__SCREAMING_SNAKE_CASE     )

        snake_case__    							=				torch.tensor([-0.6113, 0.1685, -0.0492]     ).to(__SCREAMING_SNAKE_CASE     )
        self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,					__SCREAMING_SNAKE_CASE   ,					atol=1e-4     )     )

