






import collections
import inspect
import unittest

from transformers import SwinvaConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
   import torch
   from torch import nn

   from transformers import SwinvaForImageClassification, SwinvaForMaskedImageModeling, SwinvaModel
   from transformers.models.swinva.modeling_swinva import SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor



class 				_A			:
    """simple docstring"""

    def __init__(				self      :     Dict ,     __SCREAMING_SNAKE_CASE      :     str ,     __SCREAMING_SNAKE_CASE      :     List[Any]=13 ,     __SCREAMING_SNAKE_CASE      :     Dict=32 ,     __SCREAMING_SNAKE_CASE      :     Optional[int]=2 ,     __SCREAMING_SNAKE_CASE      :     Optional[int]=3 ,     __SCREAMING_SNAKE_CASE      :     List[str]=16 ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any]=[1, 2, 1] ,     __SCREAMING_SNAKE_CASE      :     List[Any]=[2, 2, 4] ,     __SCREAMING_SNAKE_CASE      :     str=2 ,     __SCREAMING_SNAKE_CASE      :     Any=2.0 ,     __SCREAMING_SNAKE_CASE      :     Optional[int]=True ,     __SCREAMING_SNAKE_CASE      :     int=0.0 ,     __SCREAMING_SNAKE_CASE      :     Dict=0.0 ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any]=0.1 ,     __SCREAMING_SNAKE_CASE      :     Any="gelu" ,     __SCREAMING_SNAKE_CASE      :     Optional[Any]=False ,     __SCREAMING_SNAKE_CASE      :     str=True ,     __SCREAMING_SNAKE_CASE      :     List[Any]=0.02 ,     __SCREAMING_SNAKE_CASE      :     Tuple=1e-5 ,     __SCREAMING_SNAKE_CASE      :     Optional[Any]=True ,     __SCREAMING_SNAKE_CASE      :     Optional[int]=None ,     __SCREAMING_SNAKE_CASE      :     str=True ,     __SCREAMING_SNAKE_CASE      :     Any=10 ,     __SCREAMING_SNAKE_CASE      :     Dict=8 ,     )      ->							List[Any]:
           __UpperCAmelCase     				=parent
           __UpperCAmelCase     				=batch_size
           __UpperCAmelCase     				=image_size
           __UpperCAmelCase     				=patch_size
           __UpperCAmelCase     				=num_channels
           __UpperCAmelCase     				=embed_dim
           __UpperCAmelCase     				=depths
           __UpperCAmelCase     				=num_heads
           __UpperCAmelCase     				=window_size
           __UpperCAmelCase     				=mlp_ratio
           __UpperCAmelCase     				=qkv_bias
           __UpperCAmelCase     				=hidden_dropout_prob
           __UpperCAmelCase     				=attention_probs_dropout_prob
           __UpperCAmelCase     				=drop_path_rate
           __UpperCAmelCase     				=hidden_act
           __UpperCAmelCase     				=use_absolute_embeddings
           __UpperCAmelCase     				=patch_norm
           __UpperCAmelCase     				=layer_norm_eps
           __UpperCAmelCase     				=initializer_range
           __UpperCAmelCase     				=is_training
           __UpperCAmelCase     				=scope
           __UpperCAmelCase     				=use_labels
           __UpperCAmelCase     				=type_sequence_label_size
           __UpperCAmelCase     				=encoder_stride



    def _a							(				self      :     Tuple      )      ->							Optional[int]:
           __UpperCAmelCase     				=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

           __UpperCAmelCase     				=None
           if self.use_labels:
                  __UpperCAmelCase     				=ids_tensor([self.batch_size] ,     self.type_sequence_label_size      )

           __UpperCAmelCase     				=self.get_config()

           return config, pixel_values, labels



    def _a							(				self      :     List[Any]      )      ->							Optional[Any]:
           return SwinvaConfig(
               image_size=self.image_size ,     patch_size=self.patch_size ,     num_channels=self.num_channels ,     embed_dim=self.embed_dim ,     depths=self.depths ,     num_heads=self.num_heads ,     window_size=self.window_size ,     mlp_ratio=self.mlp_ratio ,     qkv_bias=self.qkv_bias ,     hidden_dropout_prob=self.hidden_dropout_prob ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob ,     drop_path_rate=self.drop_path_rate ,     hidden_act=self.hidden_act ,     use_absolute_embeddings=self.use_absolute_embeddings ,     path_norm=self.patch_norm ,     layer_norm_eps=self.layer_norm_eps ,     initializer_range=self.initializer_range ,     encoder_stride=self.encoder_stride ,     )



    def _a							(				self      :     int ,     __SCREAMING_SNAKE_CASE      :     Any ,     __SCREAMING_SNAKE_CASE      :     Tuple ,     __SCREAMING_SNAKE_CASE      :     Dict      )      ->							Optional[int]:
           __UpperCAmelCase     				=SwinvaModel(config=__SCREAMING_SNAKE_CASE      )
           model.to(__SCREAMING_SNAKE_CASE      )
           model.eval()
           __UpperCAmelCase     				=model(__SCREAMING_SNAKE_CASE      )

           __UpperCAmelCase     				=((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths      ) - 1))
           __UpperCAmelCase     				=int(config.embed_dim * 2 ** (len(config.depths      ) - 1)      )

           self.parent.assertEqual(result.last_hidden_state.shape ,     (self.batch_size, expected_seq_len, expected_dim)      )



    def _a							(				self      :     int ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any] ,     __SCREAMING_SNAKE_CASE      :     List[Any] ,     __SCREAMING_SNAKE_CASE      :     int      )      ->							Tuple:
           __UpperCAmelCase     				=SwinvaForMaskedImageModeling(config=__SCREAMING_SNAKE_CASE      )
           model.to(__SCREAMING_SNAKE_CASE      )
           model.eval()
           __UpperCAmelCase     				=model(__SCREAMING_SNAKE_CASE      )
           self.parent.assertEqual(
               result.logits.shape ,     (self.batch_size, self.num_channels, self.image_size, self.image_size)      )

           # test greyscale images
           __UpperCAmelCase     				=1
           __UpperCAmelCase     				=SwinvaForMaskedImageModeling(__SCREAMING_SNAKE_CASE      )
           model.to(__SCREAMING_SNAKE_CASE      )
           model.eval()

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



    def _a							(				self      :     Optional[int] ,     __SCREAMING_SNAKE_CASE      :     str ,     __SCREAMING_SNAKE_CASE      :     Optional[int] ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any]      )      ->							Tuple:
           __UpperCAmelCase     				=self.type_sequence_label_size
           __UpperCAmelCase     				=SwinvaForImageClassification(__SCREAMING_SNAKE_CASE      )
           model.to(__SCREAMING_SNAKE_CASE      )
           model.eval()
           __UpperCAmelCase     				=model(__SCREAMING_SNAKE_CASE ,     labels=__SCREAMING_SNAKE_CASE      )
           self.parent.assertEqual(result.logits.shape ,     (self.batch_size, self.type_sequence_label_size)      )



    def _a							(				self      :     List[str]      )      ->							Tuple:
           __UpperCAmelCase     				=self.prepare_config_and_inputs()
           __UpperCAmelCase				, __UpperCAmelCase				, __UpperCAmelCase     				=config_and_inputs
           __UpperCAmelCase     				={"""pixel_values""": pixel_values}
           return config, inputs_dict




@require_torch
class 				_A			(     UpperCamelCase					,					UpperCamelCase					,					unittest.TestCase				):
    """simple docstring"""

    lowerCamelCase	:							Optional[int]				= (
        (SwinvaModel, SwinvaForImageClassification, SwinvaForMaskedImageModeling) if is_torch_available() else ()
    )
    lowerCamelCase	:							Tuple				= (
        {'feature-extraction': SwinvaModel, 'image-classification': SwinvaForImageClassification}
        if is_torch_available()
        else {}
    )

    lowerCamelCase	:							Dict				= False
    lowerCamelCase	:							Tuple				= False
    lowerCamelCase	:							List[str]				= False
    lowerCamelCase	:							Tuple				= False



    def _a							(				self      :     str      )      ->							str:
           __UpperCAmelCase     				=SwinvaModelTester(self      )
           __UpperCAmelCase     				=ConfigTester(self ,     config_class=__SCREAMING_SNAKE_CASE ,     embed_dim=37      )



    def _a							(				self      :     List[Any]      )      ->							Optional[int]:
           self.config_tester.create_and_test_config_to_json_string()
           self.config_tester.create_and_test_config_to_json_file()
           self.config_tester.create_and_test_config_from_and_save_pretrained()
           self.config_tester.create_and_test_config_with_num_labels()
           self.config_tester.check_config_can_be_init_without_params()
           self.config_tester.check_config_arguments_init()



    def _a							(				self      :     str      )      ->							str:
           __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE      )



    @unittest.skip(reason="""Got `CUDA error: misaligned address` with PyTorch 2.0.0."""      )
    def _a							(				self      :     Tuple      )      ->							Tuple:
           pass



    @unittest.skip(reason="""Swinv2 does not use inputs_embeds"""      )
    def _a							(				self      :     Optional[Any]      )      ->							int:
           pass



    def _a							(				self      :     Tuple      )      ->							int:
           __UpperCAmelCase				, __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  __UpperCAmelCase     				=model_class(__SCREAMING_SNAKE_CASE      )
                  self.assertIsInstance(model.get_input_embeddings() ,     (nn.Module)      )
                  __UpperCAmelCase     				=model.get_output_embeddings()
                  self.assertTrue(x is None or isinstance(__SCREAMING_SNAKE_CASE ,     nn.Linear      )      )



    def _a							(				self      :     str      )      ->							List[str]:
           __UpperCAmelCase				, __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs_for_common()

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

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



    def _a							(				self      :     Tuple      )      ->							Tuple:
           __UpperCAmelCase				, __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs_for_common()
           __UpperCAmelCase     				=True

           for model_class in self.all_model_classes:
                  __UpperCAmelCase     				=True
                  __UpperCAmelCase     				=False
                  __UpperCAmelCase     				=True
                  __UpperCAmelCase     				=model_class(__SCREAMING_SNAKE_CASE      )
                  model.to(__SCREAMING_SNAKE_CASE      )
                  model.eval()
                  with torch.no_grad():
                         __UpperCAmelCase     				=model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      )      )
                  __UpperCAmelCase     				=outputs.attentions
                  __UpperCAmelCase     				=len(self.model_tester.depths      )
                  self.assertEqual(len(__SCREAMING_SNAKE_CASE      ) ,     __SCREAMING_SNAKE_CASE      )

                  # check that output_attentions also work using config
                  del inputs_dict["output_attentions"]
                  __UpperCAmelCase     				=True
                  __UpperCAmelCase     				=config.window_size**2
                  __UpperCAmelCase     				=model_class(__SCREAMING_SNAKE_CASE      )
                  model.to(__SCREAMING_SNAKE_CASE      )
                  model.eval()
                  with torch.no_grad():
                         __UpperCAmelCase     				=model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      )      )
                  __UpperCAmelCase     				=outputs.attentions
                  self.assertEqual(len(__SCREAMING_SNAKE_CASE      ) ,     __SCREAMING_SNAKE_CASE      )

                  self.assertListEqual(
                      list(attentions[0].shape[-3:]      ) ,     [self.model_tester.num_heads[0], window_size_squared, window_size_squared] ,     )
                  __UpperCAmelCase     				=len(__SCREAMING_SNAKE_CASE      )

                  # Check attention is always last and order is fine
                  __UpperCAmelCase     				=True
                  __UpperCAmelCase     				=True
                  __UpperCAmelCase     				=model_class(__SCREAMING_SNAKE_CASE      )
                  model.to(__SCREAMING_SNAKE_CASE      )
                  model.eval()
                  with torch.no_grad():
                         __UpperCAmelCase     				=model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      )      )

                  if hasattr(self.model_tester ,     """num_hidden_states_types"""      ):
                         __UpperCAmelCase     				=self.model_tester.num_hidden_states_types
                  else:
                         # also another +1 for reshaped_hidden_states
                         __UpperCAmelCase     				=2
                  self.assertEqual(out_len + added_hidden_states ,     len(__SCREAMING_SNAKE_CASE      )      )

                  __UpperCAmelCase     				=outputs.attentions

                  self.assertEqual(len(__SCREAMING_SNAKE_CASE      ) ,     __SCREAMING_SNAKE_CASE      )

                  self.assertListEqual(
                      list(self_attentions[0].shape[-3:]      ) ,     [self.model_tester.num_heads[0], window_size_squared, window_size_squared] ,     )



    def _a							(				self      :     Optional[int] ,     __SCREAMING_SNAKE_CASE      :     int ,     __SCREAMING_SNAKE_CASE      :     Optional[Any] ,     __SCREAMING_SNAKE_CASE      :     List[str] ,     __SCREAMING_SNAKE_CASE      :     Optional[Any]      )      ->							int:
           __UpperCAmelCase     				=model_class(__SCREAMING_SNAKE_CASE      )
           model.to(__SCREAMING_SNAKE_CASE      )
           model.eval()

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

           __UpperCAmelCase     				=outputs.hidden_states

           __UpperCAmelCase     				=getattr(
               self.model_tester ,     """expected_num_hidden_layers""" ,     len(self.model_tester.depths      ) + 1      )
           self.assertEqual(len(__SCREAMING_SNAKE_CASE      ) ,     __SCREAMING_SNAKE_CASE      )

           # Swinv2 has a different seq_length
           __UpperCAmelCase     				=(
               config.patch_size
               if isinstance(config.patch_size ,     collections.abc.Iterable      )
               else (config.patch_size, config.patch_size)
           )

           __UpperCAmelCase     				=(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

           self.assertListEqual(
               list(hidden_states[0].shape[-2:]      ) ,     [num_patches, self.model_tester.embed_dim] ,     )

           __UpperCAmelCase     				=outputs.reshaped_hidden_states
           self.assertEqual(len(__SCREAMING_SNAKE_CASE      ) ,     __SCREAMING_SNAKE_CASE      )

           __UpperCAmelCase				, __UpperCAmelCase				, __UpperCAmelCase				, __UpperCAmelCase     				=reshaped_hidden_states[0].shape
           __UpperCAmelCase     				=(
               reshaped_hidden_states[0].view(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     height * width      ).permute(0 ,     2 ,     1      )
           )
           self.assertListEqual(
               list(reshaped_hidden_states.shape[-2:]      ) ,     [num_patches, self.model_tester.embed_dim] ,     )



    def _a							(				self      :     str      )      ->							Union[str, Any]:
           __UpperCAmelCase				, __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs_for_common()

           __UpperCAmelCase     				=(
               self.model_tester.image_size
               if isinstance(self.model_tester.image_size ,     collections.abc.Iterable      )
               else (self.model_tester.image_size, self.model_tester.image_size)
           )

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

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

                  self.check_hidden_states_output(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      )



    def _a							(				self      :     Optional[int]      )      ->							Tuple:
           __UpperCAmelCase				, __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs_for_common()
           __UpperCAmelCase     				=3

           __UpperCAmelCase     				=(
               self.model_tester.image_size
               if isinstance(self.model_tester.image_size ,     collections.abc.Iterable      )
               else (self.model_tester.image_size, self.model_tester.image_size)
           )
           __UpperCAmelCase     				=(
               config.patch_size
               if isinstance(config.patch_size ,     collections.abc.Iterable      )
               else (config.patch_size, config.patch_size)
           )

           __UpperCAmelCase     				=image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
           __UpperCAmelCase     				=image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

           for model_class in self.all_model_classes:
                  __UpperCAmelCase     				=True
                  self.check_hidden_states_output(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     (padded_height, padded_width)      )

                  # check that output_hidden_states also work using config
                  del inputs_dict["output_hidden_states"]
                  __UpperCAmelCase     				=True
                  self.check_hidden_states_output(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     (padded_height, padded_width)      )



    def _a							(				self      :     Optional[int]      )      ->							Tuple:
           __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_image_modeling(*__SCREAMING_SNAKE_CASE      )



    def _a							(				self      :     Tuple      )      ->							Dict:
           __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE      )



    @slow
    def _a							(				self      :     int      )      ->							Dict:
           for model_name in SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  __UpperCAmelCase     				=SwinvaModel.from_pretrained(__SCREAMING_SNAKE_CASE      )
                  self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )



    def _a							(				self      :     Dict      )      ->							Union[str, Any]:
           __UpperCAmelCase				, __UpperCAmelCase     				=self.model_tester.prepare_config_and_inputs_for_common()

           __UpperCAmelCase     				=_config_zero_init(__SCREAMING_SNAKE_CASE      )
           for model_class in self.all_model_classes:
                  __UpperCAmelCase     				=model_class(config=__SCREAMING_SNAKE_CASE      )
                  for name, param in model.named_parameters():
                         if "embeddings" not in name and "logit_scale" not in name and param.requires_grad:
                                self.assertIn(
                                    ((param.data.mean() * 1e9).round() / 1e9).item() ,     [0.0, 1.0] ,     msg=f'''Parameter {name} of model {model_class} seems not properly initialized''' ,     )




@require_vision
@require_torch
class 				_A			(     unittest.TestCase				):
    """simple docstring"""

    @cached_property
    def _a							(				self      :     Tuple      )      ->							Dict:
           return (
               AutoImageProcessor.from_pretrained("""microsoft/swinv2-tiny-patch4-window8-256"""      )
               if is_vision_available()
               else None
           )



    @slow
    def _a							(				self      :     int      )      ->							Optional[int]:
           __UpperCAmelCase     				=SwinvaForImageClassification.from_pretrained("""microsoft/swinv2-tiny-patch4-window8-256"""      ).to(
               __SCREAMING_SNAKE_CASE      )
           __UpperCAmelCase     				=self.default_image_processor

           __UpperCAmelCase     				=Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
           __UpperCAmelCase     				=image_processor(images=__SCREAMING_SNAKE_CASE ,     return_tensors="""pt"""      ).to(__SCREAMING_SNAKE_CASE      )

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

           # verify the logits
           __UpperCAmelCase     				=torch.Size((1, 1000)      )
           self.assertEqual(outputs.logits.shape ,     __SCREAMING_SNAKE_CASE      )
           __UpperCAmelCase     				=torch.tensor([-0.3_947, -0.4_306, 0.0_026]      ).to(__SCREAMING_SNAKE_CASE      )
           self.assertTrue(torch.allclose(outputs.logits[0, :3] ,     __SCREAMING_SNAKE_CASE ,     atol=1e-4      )      )









'''simple docstring'''





import baseaa
import io
import json
import os
from copy import deepcopy

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler







class        UpperCAmelCase_			:
    """simple docstring"""







    def __init__(       self   :	Tuple					,	snake_case_   :	List[str]		):
          if isinstance(snake_case_					,	snake_case_		):
                # Don't modify user's data should they want to reuse it (e.g. in tests), because once we
                # modified it, it will not be accepted here again, since `auto` values would have been overridden
                snake_case__    :       Optional[Any]       		= deepcopy(snake_case_		)
          elif os.path.exists(snake_case_		):
                with io.open(snake_case_					,	"""r"""					,	encoding="""utf-8"""		) as f:
                      snake_case__    :       Union[str, Any]       		= json.load(snake_case_		)
          else:
                try:
                      snake_case__    :       List[Any]       		= baseaa.urlsafe_baadecode(snake_case_		).decode("""utf-8"""		)
                      snake_case__    :       Optional[int]       		= json.loads(snake_case_		)
                except (UnicodeDecodeError, AttributeError, ValueError):
                      raise ValueError(
                          f"Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}"		)

          snake_case__    :       str       		= config

          self.set_stage_and_offload()




    def lowerCamelCase			(       self   :	int		):
          # zero stage - this is done as early as possible, before model is created, to allow
          # ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
          # during ``zero.Init()`` which needs to know the dtype, and some other hparams.
          snake_case__    :       Union[str, Any]       		= self.get_value("""zero_optimization.stage"""					,	-1		)

          # offload
          snake_case__    :       Optional[Any]       		= False
          if self.is_zeroa() or self.is_zeroa():
                snake_case__    :       int       		= set(["""cpu""", """nvme"""]		)
                snake_case__    :       Any       		= set(
                    [
                        self.get_value("""zero_optimization.offload_optimizer.device"""		),
                        self.get_value("""zero_optimization.offload_param.device"""		),
                    ]		)
                if len(offload_devices & offload_devices_valid		) > 0:
                      snake_case__    :       List[Any]       		= True




    def lowerCamelCase			(       self   :	Any					,	snake_case_   :	Optional[Any]		):
          snake_case__    :       Optional[Any]       		= self.config

          # find the config node of interest if it exists
          snake_case__    :       List[Any]       		= ds_key_long.split("""."""		)
          snake_case__    :       Tuple       		= nodes.pop()
          for node in nodes:
                snake_case__    :       Optional[Any]       		= config.get(snake_case_		)
                if config is None:
                      return None, ds_key

          return config, ds_key




    def lowerCamelCase			(       self   :	List[Any]					,	snake_case_   :	int					,	snake_case_   :	List[str]=None		):
          snake_case__		,							snake_case__    :       Dict       		= self.find_config_node(snake_case_		)
          if config is None:
                return default
          return config.get(snake_case_					,	snake_case_		)




    def lowerCamelCase			(       self   :	Dict					,	snake_case_   :	Optional[int]					,	snake_case_   :	int=False		):
          snake_case__    :       Dict       		= self.config

          # find the config node of interest if it exists
          snake_case__    :       str       		= ds_key_long.split("""."""		)
          for node in nodes:
                snake_case__    :       Union[str, Any]       		= config
                snake_case__    :       List[Any]       		= config.get(snake_case_		)
                if config is None:
                      if must_exist:
                            raise ValueError(f"Can't find {ds_key_long} entry in the config: {self.config}"		)
                      else:
                            return

        # if found remove it
          if parent_config is not None:
                parent_config.pop(snake_case_		)




    def lowerCamelCase			(       self   :	str					,	snake_case_   :	Optional[Any]		):
          snake_case__    :       Dict       		= self.get_value(snake_case_		)
          return False if value is None else bool(snake_case_		)




    def lowerCamelCase			(       self   :	List[str]					,	snake_case_   :	Union[str, Any]		):
          snake_case__    :       Optional[int]       		= self.get_value(snake_case_		)
          return False if value is None else not bool(snake_case_		)




    def lowerCamelCase			(       self   :	Optional[Any]		):
          return self._stage == 2




    def lowerCamelCase			(       self   :	Dict		):
          return self._stage == 3




    def lowerCamelCase			(       self   :	List[Any]		):
          return self._offload







class        UpperCAmelCase_			:
    """simple docstring"""







    def __init__(       self   :	Any					,	snake_case_   :	List[Any]		):
          snake_case__    :       Tuple       		= engine




    def lowerCamelCase			(       self   :	Dict					,	snake_case_   :	Optional[Any]					,	**snake_case_   :	int		):
          # runs backpropagation and handles mixed precision
          self.engine.backward(snake_case_					,	**snake_case_		)

          # Deepspeed's `engine.step` performs the following operations:
          # - gradient accumulation check
          # - gradient clipping
          # - optimizer step
          # - zero grad
          # - checking overflow
          # - lr_scheduler step (only if engine.lr_scheduler is not None)
          self.engine.step()
          # and this plugin overrides the above calls with no-ops when Accelerate runs under
          # Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
          # training loop that works transparently under many training regimes.







class        UpperCAmelCase_			(			_a      ):
    """simple docstring"""







    def __init__(       self   :	Tuple					,	snake_case_   :	Any		):
          super().__init__(snake_case_					,	device_placement=snake_case_					,	scaler=snake_case_		)
          snake_case__    :       Optional[int]       		= hasattr(self.optimizer					,	"""overflow"""		)




    def lowerCamelCase			(       self   :	int					,	snake_case_   :	Optional[Any]=None		):
          pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




    def lowerCamelCase			(       self   :	List[Any]		):
          pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




    @property
    def lowerCamelCase			(       self   :	List[str]		):
          if self.__has_overflow__:
                return self.optimizer.overflow
          return False







class        UpperCAmelCase_			(			_a      ):
    """simple docstring"""







    def __init__(       self   :	Tuple					,	snake_case_   :	Optional[Any]					,	snake_case_   :	Tuple		):
          super().__init__(snake_case_					,	snake_case_		)




    def lowerCamelCase			(       self   :	Optional[int]		):
          pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







class        UpperCAmelCase_			:
    """simple docstring"""







    def __init__(       self   :	Optional[int]					,	snake_case_   :	Dict					,	snake_case_   :	int=0.001					,	snake_case_   :	Optional[Any]=0					,	**snake_case_   :	List[Any]		):
          snake_case__    :       int       		= params
          snake_case__    :       Optional[Any]       		= lr
          snake_case__    :       Any       		= weight_decay
          snake_case__    :       Optional[Any]       		= kwargs







class        UpperCAmelCase_			:
    """simple docstring"""







    def __init__(       self   :	Dict					,	snake_case_   :	str					,	snake_case_   :	Optional[int]=None					,	snake_case_   :	Any=0					,	**snake_case_   :	List[str]		):
          snake_case__    :       List[Any]       		= optimizer
          snake_case__    :       List[Any]       		= total_num_steps
          snake_case__    :       List[Any]       		= warmup_num_steps
          snake_case__    :       Optional[Any]       		= kwargs
