





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   :							Dict     =	get_logger(__name__)
def 				__UpperCamelCase  (		snake_case    ,	snake_case    ,	snake_case    ,	snake_case    ,	snake_case    ,	snake_case = None    ,	)  ->						Optional[int]:


   '''simple docstring'''







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


   '''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:
      __A				 =       '''cpu'''
   elif map_location == "on_device":
      __A				 =       PartialState().device
   # Model states
   for i, model in enumerate(snake_case     ):
      __A				 =       F"{MODEL_NAME}.bin" if i == 0 else F"{MODEL_NAME}_{i}.bin"
      __A				 =       os.path.join(snake_case    ,	snake_case     )
      models[i].load_state_dict(torch.load(snake_case    ,	map_location=snake_case     )    ,	**snake_case     )
   logger.info('''All model weights loaded successfully'''     )

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

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

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

   # Random states
   try:
      __A				 =       torch.load(os.path.join(snake_case    ,	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 				__UpperCamelCase  (		snake_case    ,	snake_case    ,	snake_case = 0     )  ->						List[str]:


   '''simple docstring'''







   __A				 =       Path(snake_case     ) / F"custom_checkpoint_{index}.pkl"
   logger.info(F"Saving the state of {get_pretty_name(snake_case     )} to {save_location}"     )
   torch.save(obj.state_dict()    ,	snake_case     )
def 				__UpperCamelCase  (		snake_case    ,	snake_case    ,	snake_case = 0     )  ->						Tuple:


   '''simple docstring'''







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








from __future__ import annotations
def 				__UpperCamelCase  (		snake_case    ,	snake_case     )  ->						list:


   '''simple docstring'''







   __A				 =       [None] * (len(snake_case     ) + len(snake_case     ))

   __A				 =       0  # pointer to current index for left Half
   __A				 =       0  # pointer to current index for the right Half
   __A				 =       0  # pointer to current index for the sorted array Half

   while pointera < len(snake_case     ) and pointera < len(snake_case     ):
      if left_half[pointera] < right_half[pointera]:
         __A				 =       left_half[pointera]
         pointera += 1
         index += 1
      else:
         __A				 =       right_half[pointera]
         pointera += 1
         index += 1
   while pointera < len(snake_case     ):
      __A				 =       left_half[pointera]
      pointera += 1
      index += 1

   while pointera < len(snake_case     ):
      __A				 =       right_half[pointera]
      pointera += 1
      index += 1

   return sorted_array
def 				__UpperCamelCase  (		snake_case     )  ->						list:


   '''simple docstring'''







   if len(snake_case     ) <= 1:
      return array
   # the actual formula to calculate the middle element = left + (right - left) // 2
   # this avoids integer overflow in case of large N
   __A				 =       0 + (len(snake_case     ) - 0) // 2

   # Split the array into halves till the array length becomes equal to One
   # merge the arrays of single length returned by mergeSort function and
   # pass them into the merge arrays function which merges the array
   __A				 =       array[:middle]
   __A				 =       array[middle:]

   return merge(merge_sort(snake_case     )    ,	merge_sort(snake_case     )     )


if __name__ == "__main__":
    import doctest

    doctest.testmod()

