

'''simple docstring'''



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
  import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


UpperCAmelCase				: Tuple         =  get_logger(__name__)







def      a__ (						a__   ,					a__   ,					a__   ,					a__   ,					a__   ,					a__ = None   ,					):




     """simple docstring"""


     for i, state in enumerate(a__		):
          __SCREAMING_SNAKE_CASE						     =   F'{MODEL_NAME}.bin' if i == 0 else F'{MODEL_NAME}_{i}.bin'
          __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
          save(a__   ,					a__		)
          logger.info(F'Model weights saved in {output_model_file}'		)
     # Optimizer states
     for i, opt in enumerate(a__		):
          __SCREAMING_SNAKE_CASE						     =   opt.state_dict()
          __SCREAMING_SNAKE_CASE						     =   F'{OPTIMIZER_NAME}.bin' if i == 0 else F'{OPTIMIZER_NAME}_{i}.bin'
          __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
          save(a__   ,					a__		)
          logger.info(F'Optimizer state saved in {output_optimizer_file}'		)
     # Scheduler states
     for i, scheduler in enumerate(a__		):
          __SCREAMING_SNAKE_CASE						     =   scheduler.state_dict()
          __SCREAMING_SNAKE_CASE						     =   F'{SCHEDULER_NAME}.bin' if i == 0 else F'{SCHEDULER_NAME}_{i}.bin'
          __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
          save(a__   ,					a__		)
          logger.info(F'Scheduler state saved in {output_scheduler_file}'		)
     # GradScaler state
     if scaler is not None:
          __SCREAMING_SNAKE_CASE						     =   scaler.state_dict()
          __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
          torch.save(a__   ,					a__		)
          logger.info(F'Gradient scaler state saved in {output_scaler_file}'		)
     # Random number generator states
     __SCREAMING_SNAKE_CASE						     =   {}
     __SCREAMING_SNAKE_CASE						     =   F'{RNG_STATE_NAME}_{process_index}.pkl'
     __SCREAMING_SNAKE_CASE						     =   random.getstate()
     __SCREAMING_SNAKE_CASE						     =   np.random.get_state()
     __SCREAMING_SNAKE_CASE						     =   torch.get_rng_state()
     if is_xpu_available():
          __SCREAMING_SNAKE_CASE						     =   torch.xpu.get_rng_state_all()
     else:
          __SCREAMING_SNAKE_CASE						     =   torch.cuda.get_rng_state_all()
     if is_tpu_available():
          __SCREAMING_SNAKE_CASE						     =   xm.get_rng_state()
     __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
     torch.save(a__   ,					a__		)
     logger.info(F'Random states saved in {output_states_file}'		)
     return output_dir







def      a__ (						a__   ,					a__   ,					a__   ,					a__   ,					a__   ,					a__=None   ,					a__=None   ,					**a__   ,					):




     """simple docstring"""


     if map_location not in [None, "cpu", "on_device"]:
          raise TypeError(
              """Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"""		)
     if map_location is None:
          __SCREAMING_SNAKE_CASE						     =   """cpu"""
     elif map_location == "on_device":
          __SCREAMING_SNAKE_CASE						     =   PartialState().device
     # Model states
     for i, model in enumerate(a__		):
          __SCREAMING_SNAKE_CASE						     =   F'{MODEL_NAME}.bin' if i == 0 else F'{MODEL_NAME}_{i}.bin'
          __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
          models[i].load_state_dict(torch.load(a__   ,					map_location=a__		)   ,					**a__		)
     logger.info("""All model weights loaded successfully"""		)

     # Optimizer states
     for i, opt in enumerate(a__		):
          __SCREAMING_SNAKE_CASE						     =   F'{OPTIMIZER_NAME}.bin' if i == 0 else F'{OPTIMIZER_NAME}_{i}.bin'
          __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
          __SCREAMING_SNAKE_CASE						     =   torch.load(a__   ,					map_location=a__		)
          optimizers[i].load_state_dict(a__		)
     logger.info("""All optimizer states loaded successfully"""		)

     # Scheduler states
     for i, scheduler in enumerate(a__		):
          __SCREAMING_SNAKE_CASE						     =   F'{SCHEDULER_NAME}.bin' if i == 0 else F'{SCHEDULER_NAME}_{i}.bin'
          __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
          scheduler.load_state_dict(torch.load(a__		)		)
     logger.info("""All scheduler states loaded successfully"""		)

     # GradScaler state
     if scaler is not None:
          __SCREAMING_SNAKE_CASE						     =   os.path.join(a__   ,					a__		)
          scaler.load_state_dict(torch.load(a__		)		)
          logger.info("""GradScaler state loaded successfully"""		)

     # Random states
     try:
          __SCREAMING_SNAKE_CASE						     =   torch.load(os.path.join(a__   ,					F'{RNG_STATE_NAME}_{process_index}.pkl'		)		)
          random.setstate(states["""random_state"""]		)
          np.random.set_state(states["""numpy_random_seed"""]		)
          torch.set_rng_state(states["""torch_manual_seed"""]		)
          if is_xpu_available():
               torch.xpu.set_rng_state_all(states["""torch_xpu_manual_seed"""]		)
          else:
               torch.cuda.set_rng_state_all(states["""torch_cuda_manual_seed"""]		)
          if is_tpu_available():
               xm.set_rng_state(states["""xm_seed"""]		)
          logger.info("""All random states loaded successfully"""		)
     except Exception:
          logger.info("""Could not load random states"""		)







def      a__ (						a__   ,					a__   ,					a__ = 0		):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE						     =   Path(a__		) / F'custom_checkpoint_{index}.pkl'
     logger.info(F'Saving the state of {get_pretty_name(a__		)} to {save_location}'		)
     torch.save(obj.state_dict()   ,					a__		)







def      a__ (						a__   ,					a__   ,					a__ = 0		):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE						     =   F'{path}/custom_checkpoint_{index}.pkl'
     logger.info(F'Loading the state of {get_pretty_name(a__		)} from {load_location}'		)
     obj.load_state_dict(torch.load(a__   ,					map_location="""cpu"""		)		)









'''simple docstring'''



import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


if is_tpu_available(check_device=False):
  import torch_xla.core.xla_model as xm







def      a__ (						a__   ,					a__		):




     """simple docstring"""


     if isinstance(a__   ,					(list, tuple)		):
          return honor_type(a__   ,					(move_to_device(a__   ,					a__		) for t in state)		)
     elif isinstance(a__   ,					a__		):
          return type(a__		)({k: move_to_device(a__   ,					a__		) for k, v in state.items()}		)
     elif isinstance(a__   ,					torch.Tensor		):
          return state.to(a__		)
     return state







class 		lowerCAmelCase__     (   torch.optim.Optimizer      ):



 """simple docstring"""



 def __init__(	self      : Tuple      ,							__SCREAMING_SNAKE_CASE      : Union[str, Any]      ,							__SCREAMING_SNAKE_CASE      : Dict=True      ,							__SCREAMING_SNAKE_CASE      : List[Any]=None    )						->				str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE						     =   optimizer
      __SCREAMING_SNAKE_CASE						     =   scaler
      __SCREAMING_SNAKE_CASE						     =   AcceleratorState()
      __SCREAMING_SNAKE_CASE						     =   GradientState()
      __SCREAMING_SNAKE_CASE						     =   device_placement
      __SCREAMING_SNAKE_CASE						     =   False
      __SCREAMING_SNAKE_CASE						     =   None

      # Handle device placement
      if device_placement:
           __SCREAMING_SNAKE_CASE						     =   self.optimizer.state_dict()
           if self.accelerator_state.distributed_type == DistributedType.TPU:
                xm.send_cpu_data_to_device(__SCREAMING_SNAKE_CASE      ,							self.accelerator_state.device    )
           else:
                __SCREAMING_SNAKE_CASE						     =   move_to_device(__SCREAMING_SNAKE_CASE      ,							self.accelerator_state.device    )
           self.optimizer.load_state_dict(__SCREAMING_SNAKE_CASE    )
 @property
 def    UpperCAmelCase__	(	self      : str    )						->				Union[str, Any]:


      """simple docstring"""




      return self.optimizer.state
 @state.setter
 def    UpperCAmelCase__	(	self      : List[str]      ,							__SCREAMING_SNAKE_CASE      : Optional[Any]    )						->				List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE						     =   state
 @property
 def    UpperCAmelCase__	(	self      : Optional[int]    )						->				str:


      """simple docstring"""




      return self.optimizer.param_groups
 @param_groups.setter
 def    UpperCAmelCase__	(	self      : Dict      ,							__SCREAMING_SNAKE_CASE      : str    )						->				List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE						     =   param_groups
 @property
 def    UpperCAmelCase__	(	self      : List[Any]    )						->				Any:


      """simple docstring"""




      return self.optimizer.defaults
 @defaults.setter
 def    UpperCAmelCase__	(	self      : Any      ,							__SCREAMING_SNAKE_CASE      : Dict    )						->				Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE						     =   defaults
 def    UpperCAmelCase__	(	self      : Any      ,							__SCREAMING_SNAKE_CASE      : Any    )						->				List[str]:


      """simple docstring"""




      self.optimizer.add_param_group(__SCREAMING_SNAKE_CASE    )
 def    UpperCAmelCase__	(	self      : List[str]      ,							__SCREAMING_SNAKE_CASE      : List[str]    )						->				Dict:


      """simple docstring"""




      if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
           xm.send_cpu_data_to_device(__SCREAMING_SNAKE_CASE      ,							self.accelerator_state.device    )
      self.optimizer.load_state_dict(__SCREAMING_SNAKE_CASE    )
 def    UpperCAmelCase__	(	self      : Union[str, Any]    )						->				Union[str, Any]:


      """simple docstring"""




      return self.optimizer.state_dict()
 def    UpperCAmelCase__	(	self      : List[str]      ,							__SCREAMING_SNAKE_CASE      : List[str]=None    )						->				Any:


      """simple docstring"""




      if self.gradient_state.sync_gradients:
           __SCREAMING_SNAKE_CASE						     =   """set_to_none""" in inspect.signature(self.optimizer.zero_grad    ).parameters
           if accept_arg:
                if set_to_none is None:
                     __SCREAMING_SNAKE_CASE						     =   False
                self.optimizer.zero_grad(set_to_none=__SCREAMING_SNAKE_CASE    )
           else:
                if set_to_none is not None:
                     raise ValueError("""`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer."""    )
                self.optimizer.zero_grad()
 def    UpperCAmelCase__	(	self      : List[str]      ,							__SCREAMING_SNAKE_CASE      : List[Any]=None    )						->				Dict:


      """simple docstring"""




      if self.gradient_state.sync_gradients:
           if self.accelerator_state.distributed_type == DistributedType.TPU:
                __SCREAMING_SNAKE_CASE						     =   {"""closure""": closure} if closure is not None else {}
                xm.optimizer_step(self.optimizer      ,							optimizer_args=__SCREAMING_SNAKE_CASE    )
           elif self.scaler is not None:
                __SCREAMING_SNAKE_CASE						     =   False
                if self._last_scale is None:
                     # `get_scale` is an async operation requiring full synchronization
                     # on CPU and GPUs before finishing. As a result, we store away
                     # the prior one to reduce the call overhead
                     __SCREAMING_SNAKE_CASE						     =   self.scaler.get_scale()
                     __SCREAMING_SNAKE_CASE						     =   True
                self.scaler.step(self.optimizer      ,							__SCREAMING_SNAKE_CASE    )
                self.scaler.update()
                __SCREAMING_SNAKE_CASE						     =   self.scaler.get_scale()
                if not new_scale:
                     # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
                     __SCREAMING_SNAKE_CASE						     =   scale_after < self._last_scale
                __SCREAMING_SNAKE_CASE						     =   scale_after

           else:
                self.optimizer.step(__SCREAMING_SNAKE_CASE    )
 def    UpperCAmelCase__	(	self      : Union[str, Any]      ,							__SCREAMING_SNAKE_CASE      : Any    )						->				Any:


      """simple docstring"""




      for param_group in self.optimizer.param_groups:
           __SCREAMING_SNAKE_CASE						     =   [parameters_map.get(__SCREAMING_SNAKE_CASE      ,							__SCREAMING_SNAKE_CASE    ) for p in param_group["""params"""]]
 @property
 def    UpperCAmelCase__	(	self      : int    )						->				Any:


      """simple docstring"""




      warnings.warn(
          """The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use """
          """`optimizer.step_was_skipped` instead."""      ,							__SCREAMING_SNAKE_CASE      ,							)
      return self._is_overflow
 @property
 def    UpperCAmelCase__	(	self      : Union[str, Any]    )						->				Optional[int]:


      """simple docstring"""




      return self._is_overflow
 def __getstate__(	self      : Tuple    )						->				str:


      """simple docstring"""




      return self.__dict__.copy()


 def __setstate__(	self      : str      ,							__SCREAMING_SNAKE_CASE      : Optional[int]    )						->				Any:


      """simple docstring"""




      self.__dict__.update(__SCREAMING_SNAKE_CASE    )






