



"""simple docstring"""







import copy
import inspect

from transformers.testing_utils import require_torch, torch_device
from transformers.utils.backbone_utils import BackboneType





@require_torch
class   UpperCamelCase	:




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_							:       Tuple						       =  ()
	SCREAMING_SNAKE_CASE_							:       Optional[Any]						       =  True






	def    lowerCamelCase__					(     self   ):
			_lowercase				: Any     =   self.config_class

			# test default config
			_lowercase				: List[str]     =   config_class()
			self.assertIsNotNone(UpperCAmelCase_   )
			_lowercase				: Optional[Any]     =   ["""stem"""] + [f"""stage{idx}""" for idx in range(1				,len(config.depths   ) + 1   )]
			self.assertEqual(config.stage_names				,UpperCAmelCase_   )
			self.assertTrue(set(config.out_features   ).issubset(set(config.stage_names   )   )   )

			# Test out_features and out_indices are correctly set
			# out_features and out_indices both None
			_lowercase				: Optional[Any]     =   config_class(out_features=UpperCAmelCase_				,out_indices=UpperCAmelCase_   )
			self.assertEqual(config.out_features				,[config.stage_names[-1]]   )
			self.assertEqual(config.out_indices				,[len(config.stage_names   ) - 1]   )

			# out_features and out_indices both set
			_lowercase				: Any     =   config_class(out_features=["""stem""", """stage1"""]				,out_indices=[0, 1]   )
			self.assertEqual(config.out_features				,["""stem""", """stage1"""]   )
			self.assertEqual(config.out_indices				,[0, 1]   )

			# Only out_features set
			_lowercase				: str     =   config_class(out_features=["""stage1""", """stage3"""]   )
			self.assertEqual(config.out_features				,["""stage1""", """stage3"""]   )
			self.assertEqual(config.out_indices				,[1, 3]   )

			# Only out_indices set
			_lowercase				: Any     =   config_class(out_indices=[0, 2]   )
			self.assertEqual(config.out_features				,[config.stage_names[0], config.stage_names[2]]   )
			self.assertEqual(config.out_indices				,[0, 2]   )

			# Error raised when out_indices do not correspond to out_features
			with self.assertRaises(UpperCAmelCase_   ):
					_lowercase				: str     =   config_class(out_features=["""stage1""", """stage2"""]				,out_indices=[0, 2]   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: Tuple     =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					_lowercase				: List[str]     =   model_class(UpperCAmelCase_   )
					_lowercase				: int     =   inspect.signature(model.forward   )
					# signature.parameters is an OrderedDict => so arg_names order is deterministic
					_lowercase				: Optional[Any]     =   [*signature.parameters.keys()]
					_lowercase				: Dict     =   ["""pixel_values"""]
					self.assertListEqual(arg_names[:1]				,UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: Tuple     =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					_lowercase				: Dict     =   model_class(UpperCAmelCase_   )
					self.assertEqual(len(model.channels   )				,len(config.out_features   )   )
					_lowercase				: Union[str, Any]     =   model.num_features
					_lowercase				: Optional[Any]     =   [config.stage_names.index(UpperCAmelCase_   ) for feat in config.out_features]
					_lowercase				: List[str]     =   [num_features[idx] for idx in out_indices]
					self.assertListEqual(model.channels				,UpperCAmelCase_   )

					_lowercase				: List[Any]     =   copy.deepcopy(UpperCAmelCase_   )
					_lowercase				: List[Any]     =   None
					_lowercase				: Optional[int]     =   model_class(UpperCAmelCase_   )
					self.assertEqual(len(model.channels   )				,1   )
					self.assertListEqual(model.channels				,[num_features[-1]]   )

					_lowercase				: int     =   copy.deepcopy(UpperCAmelCase_   )
					_lowercase				: Optional[int]     =   None
					_lowercase				: List[str]     =   model_class(UpperCAmelCase_   )
					self.assertEqual(len(model.channels   )				,1   )
					self.assertListEqual(model.channels				,[num_features[-1]]   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: Optional[Any]     =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					_lowercase				: Optional[int]     =   model_class(UpperCAmelCase_   )
					model.to(UpperCAmelCase_   )
					model.eval()
					_lowercase				: Optional[int]     =   model(**UpperCAmelCase_   )

					self.assertEqual(len(result.feature_maps   )				,len(config.out_features   )   )
					self.assertEqual(len(model.channels   )				,len(config.out_features   )   )
					self.assertEqual(len(result.feature_maps   )				,len(config.out_indices   )   )
					self.assertEqual(len(model.channels   )				,len(config.out_indices   )   )

					# Check output of last stage is taken if out_features=None, out_indices=None
					_lowercase				: str     =   copy.deepcopy(UpperCAmelCase_   )
					_lowercase				: Dict     =   None
					_lowercase				: Union[str, Any]     =   model_class(UpperCAmelCase_   )
					model.to(UpperCAmelCase_   )
					model.eval()
					_lowercase				: int     =   model(**UpperCAmelCase_   )

					self.assertEqual(len(result.feature_maps   )				,1   )
					self.assertEqual(len(model.channels   )				,1   )

					_lowercase				: List[str]     =   copy.deepcopy(UpperCAmelCase_   )
					_lowercase				: Any     =   None
					_lowercase				: int     =   model_class(UpperCAmelCase_   )
					model.to(UpperCAmelCase_   )
					model.eval()
					_lowercase				: Optional[int]     =   model(**UpperCAmelCase_   )

					self.assertEqual(len(result.feature_maps   )				,1   )
					self.assertEqual(len(model.channels   )				,1   )

					# Check backbone can be initialized with fresh weights
					_lowercase				: int     =   copy.deepcopy(UpperCAmelCase_   )
					_lowercase				: int     =   False
					_lowercase				: int     =   model_class(UpperCAmelCase_   )
					model.to(UpperCAmelCase_   )
					model.eval()
					_lowercase				: Optional[int]     =   model(**UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: Optional[int]     =   self.model_tester.prepare_config_and_inputs_for_common()

			for backbone_class in self.all_model_classes:
					_lowercase				: List[Any]     =   backbone_class(UpperCAmelCase_   )

					self.assertTrue(hasattr(UpperCAmelCase_				,"""backbone_type"""   )   )
					self.assertTrue(hasattr(UpperCAmelCase_				,"""stage_names"""   )   )
					self.assertTrue(hasattr(UpperCAmelCase_				,"""num_features"""   )   )
					self.assertTrue(hasattr(UpperCAmelCase_				,"""out_indices"""   )   )
					self.assertTrue(hasattr(UpperCAmelCase_				,"""out_features"""   )   )
					self.assertTrue(hasattr(UpperCAmelCase_				,"""out_feature_channels"""   )   )
					self.assertTrue(hasattr(UpperCAmelCase_				,"""channels"""   )   )

					self.assertIsInstance(backbone.backbone_type				,UpperCAmelCase_   )
					# Verify num_features has been initialized in the backbone init
					self.assertIsNotNone(backbone.num_features   )
					self.assertTrue(len(backbone.channels   ) == len(backbone.out_indices   )   )
					self.assertTrue(len(backbone.stage_names   ) == len(backbone.num_features   )   )
					self.assertTrue(len(backbone.channels   ) <= len(backbone.num_features   )   )
					self.assertTrue(len(backbone.out_feature_channels   ) == len(backbone.stage_names   )   )






	def    lowerCamelCase__					(     self   ):
			_lowercase     ,  _lowercase				: int     =   self.model_tester.prepare_config_and_inputs_for_common()
			_lowercase				: int     =   inputs_dict["""pixel_values"""].shape[0]

			for backbone_class in self.all_model_classes:
					_lowercase				: List[str]     =   backbone_class(UpperCAmelCase_   )
					backbone.to(UpperCAmelCase_   )
					backbone.eval()

					_lowercase				: str     =   backbone(**UpperCAmelCase_   )

					# Test default outputs and verify feature maps
					self.assertIsInstance(outputs.feature_maps				,UpperCAmelCase_   )
					self.assertTrue(len(outputs.feature_maps   ) == len(backbone.channels   )   )
					for feature_map, n_channels in zip(outputs.feature_maps				,backbone.channels   ):
							self.assertTrue(feature_map.shape[:2]				,(batch_size, n_channels)   )
					self.assertIsNone(outputs.hidden_states   )
					self.assertIsNone(outputs.attentions   )

					# Test output_hidden_states=True
					_lowercase				: List[Any]     =   backbone(**UpperCAmelCase_				,output_hidden_states=UpperCAmelCase_   )
					self.assertIsNotNone(outputs.hidden_states   )
					self.assertTrue(len(outputs.hidden_states   )				,len(backbone.stage_names   )   )
					for hidden_state, n_channels in zip(outputs.hidden_states				,backbone.channels   ):
							self.assertTrue(hidden_state.shape[:2]				,(batch_size, n_channels)   )

					# Test output_attentions=True
					if self.has_attentions:
							_lowercase				: Any     =   backbone(**UpperCAmelCase_				,output_attentions=UpperCAmelCase_   )
							self.assertIsNotNone(outputs.attentions   )




'''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 				__UpperCamelCase (					__UpperCAmelCase  :					List[Any]=2 ,   __UpperCAmelCase  :					Dict=True ,   __UpperCAmelCase  :					str=False			)							->	List[Any]:
  lowercase				:							Tuple									=   Accelerator(step_scheduler_with_optimizer=__UpperCAmelCase ,   split_batches=__UpperCAmelCase			)
  lowercase				:							Optional[int]									=   torch.nn.Linear(2 ,   4			)
  lowercase				:							Dict									=   torch.optim.AdamW(model.parameters() ,   lr=1.0			)
  lowercase				:							List[Any]									=   torch.optim.lr_scheduler.OneCycleLR(__UpperCAmelCase ,   max_lr=0.01 ,   steps_per_epoch=2 ,   epochs=1			)
  lowercase				,      lowercase				,      lowercase				:							Optional[Any]									=   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 				__UpperCamelCase (					__UpperCAmelCase  :					str=2 ,   __UpperCAmelCase  :					str=True ,   __UpperCAmelCase  :					Tuple=False			)							->	str:
  lowercase				:							List[Any]									=   Accelerator(step_scheduler_with_optimizer=__UpperCAmelCase ,   split_batches=__UpperCAmelCase			)
  lowercase				:							Optional[Any]									=   torch.nn.Linear(2 ,   4			)
  lowercase				:							List[str]									=   torch.optim.AdamW(model.parameters() ,   lr=1.0			)
  lowercase				:							Dict									=   torch.optim.lr_scheduler.LambdaLR(__UpperCAmelCase ,   lr_lambda=lambda __UpperCAmelCase			: 1 - n / 10			)
  lowercase				,      lowercase				,      lowercase				:							int									=   accelerator.prepare(__UpperCAmelCase ,   __UpperCAmelCase ,   __UpperCAmelCase			)

  # Optimizer has stepped
  lowercase				:							Any									=   False
  scheduler.step()
  lowercase				:							Optional[Any]									=   1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
  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
  lowercase				:							List[Any]									=   True
  scheduler.step()
  if not step_scheduler_with_optimizer:
    lowercase				:							List[Any]									=   1 - 2 / 10
  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 				__UpperCamelCase (					__UpperCAmelCase  :					int = 2			)							->	Dict:
  from transformers import get_linear_schedule_with_warmup

  lowercase				:							Optional[Any]									=   [1, 2, 4]
  for num_steps in steps:
    lowercase				:							Dict									=   GradientAccumulationPlugin(num_steps=__UpperCAmelCase ,   adjust_scheduler=num_steps > 1			)
    lowercase				:							Tuple									=   Accelerator(gradient_accumulation_plugin=__UpperCAmelCase			)
    lowercase				:							str									=   torch.nn.Linear(2 ,   4			)
    lowercase				:							Dict									=   torch.optim.AdamW(model.parameters() ,   lr=10.0			)
    lowercase				:							List[str]									=   get_linear_schedule_with_warmup(optimizer=__UpperCAmelCase ,   num_warmup_steps=0 ,   num_training_steps=20			)

    lowercase				,      lowercase				,      lowercase				:							Dict									=   accelerator.prepare(__UpperCAmelCase ,   __UpperCAmelCase ,   __UpperCAmelCase			)

    for i in range(10 * num_steps			):
      with accelerator.accumulate(__UpperCAmelCase			):
        optimizer.step()
        scheduler.step()

      if i == (10 * 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 _lowerCAmelCase			(	unittest.TestCase					):







      def      lowerCamelCase__      (  self     ):
        debug_launcher(partial(lowercase  ,    num_processes=1     )  ,    num_processes=1     )
        debug_launcher(partial(lowercase  ,    num_processes=1  ,    split_batches=lowercase     )  ,    num_processes=1     )







      def      lowerCamelCase__      (  self     ):
        debug_launcher(partial(lowercase  ,    num_processes=1     )  ,    num_processes=1     )
        debug_launcher(partial(lowercase  ,    num_processes=1  ,    split_batches=lowercase     )  ,    num_processes=1     )







      def      lowerCamelCase__      (  self     ):
        debug_launcher(partial(lowercase  ,    num_processes=1  ,    step_scheduler_with_optimizer=lowercase     )  ,    num_processes=1     )







      def      lowerCamelCase__      (  self     ):
        debug_launcher(partial(lowercase  ,    num_processes=1  ,    step_scheduler_with_optimizer=lowercase     )  ,    num_processes=1     )







      def      lowerCamelCase__      (  self     ):
        AcceleratorState._reset_state(lowercase     )
        debug_launcher(lowercase     )
        debug_launcher(partial(lowercase  ,    num_processes=1  ,    split_batches=lowercase     )  ,    num_processes=1     )







      def      lowerCamelCase__      (  self     ):
        AcceleratorState._reset_state(lowercase     )
        debug_launcher(lowercase     )
        debug_launcher(partial(lowercase  ,    num_processes=1  ,    split_batches=lowercase     )  ,    num_processes=1     )







      def      lowerCamelCase__      (  self     ):
        AcceleratorState._reset_state(lowercase     )
        debug_launcher(partial(lowercase  ,    step_scheduler_with_optimizer=lowercase     )     )







      def      lowerCamelCase__      (  self     ):
        AcceleratorState._reset_state(lowercase     )
        debug_launcher(partial(lowercase  ,    step_scheduler_with_optimizer=lowercase     )     )







      @require_huggingface_suite
      def      lowerCamelCase__      (  self     ):
        AcceleratorState._reset_state(lowercase     )
        debug_launcher(partial(lowercase  ,    num_processes=1     )     )
        debug_launcher(lowercase     )