




'''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


_SCREAMING_SNAKE_CASE =	get_logger(__name__)





def  lowerCamelCase(       SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_ = None					,)		->		Optional[Any]:
			# Model states
			for i, state in enumerate(SCREAMING_SNAKE_CASE_    ):
						A_       		=   F'{MODEL_NAME}.bin' if i == 0 else F'{MODEL_NAME}_{i}.bin'
						A_       		=   os.path.join(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
						save(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
						logger.info(F'Model weights saved in {output_model_file}'    )
			# Optimizer states
			for i, opt in enumerate(SCREAMING_SNAKE_CASE_    ):
						A_       		=   opt.state_dict()
						A_       		=   F'{OPTIMIZER_NAME}.bin' if i == 0 else F'{OPTIMIZER_NAME}_{i}.bin'
						A_       		=   os.path.join(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
						save(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
						logger.info(F'Optimizer state saved in {output_optimizer_file}'    )
			# Scheduler states
			for i, scheduler in enumerate(SCREAMING_SNAKE_CASE_    ):
						A_       		=   scheduler.state_dict()
						A_       		=   F'{SCHEDULER_NAME}.bin' if i == 0 else F'{SCHEDULER_NAME}_{i}.bin'
						A_       		=   os.path.join(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
						save(SCREAMING_SNAKE_CASE_					,SCREAMING_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(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
						torch.save(SCREAMING_SNAKE_CASE_					,SCREAMING_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(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
			torch.save(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
			logger.info(F'Random states saved in {output_states_file}'    )
			return output_dir





def  lowerCamelCase(       SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_=None					,SCREAMING_SNAKE_CASE_=None					,**SCREAMING_SNAKE_CASE_					,)		->		str:
			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(SCREAMING_SNAKE_CASE_    ):
						A_       		=   F'{MODEL_NAME}.bin' if i == 0 else F'{MODEL_NAME}_{i}.bin'
						A_       		=   os.path.join(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )
						models[i].load_state_dict(torch.load(SCREAMING_SNAKE_CASE_					,map_location=SCREAMING_SNAKE_CASE_    )					,**SCREAMING_SNAKE_CASE_    )
			logger.info('All model weights loaded successfully'    )

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

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

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

			# Random states
			try:
						A_       		=   torch.load(os.path.join(SCREAMING_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  lowerCamelCase(       SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_ = 0    )		->		Union[str, Any]:
			A_       		=   Path(SCREAMING_SNAKE_CASE_    ) / F'custom_checkpoint_{index}.pkl'
			logger.info(F'Saving the state of {get_pretty_name(SCREAMING_SNAKE_CASE_    )} to {save_location}'    )
			torch.save(obj.state_dict()					,SCREAMING_SNAKE_CASE_    )






def  lowerCamelCase(       SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_ = 0    )		->		str:
			A_       		=   F'{path}/custom_checkpoint_{index}.pkl'
			logger.info(F'Loading the state of {get_pretty_name(SCREAMING_SNAKE_CASE_    )} from {load_location}'    )
			obj.load_state_dict(torch.load(SCREAMING_SNAKE_CASE_					,map_location='cpu'    )    )










'''simple docstring'''
import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


_SCREAMING_SNAKE_CASE =	namedtuple("ExampleNamedTuple", "a b c")
class      __UpperCAmelCase							(    unittest.TestCase							):



					'''simple docstring'''





					def 						__snake_case	(  self    :	List[Any])						->       Tuple:
								A_       		=   torch.randn(5			,					2)
								A_       		=   torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

								A_       		=   send_to_device(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta.cpu()			,					_lowercase))

								A_       		=   send_to_device((tensor, [tensor, tensor], 1)			,					_lowercase)
								self.assertIsInstance(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta[0].cpu()			,					_lowercase))
								self.assertIsInstance(resulta[1]			,					_lowercase)
								self.assertTrue(torch.equal(resulta[1][0].cpu()			,					_lowercase))
								self.assertTrue(torch.equal(resulta[1][1].cpu()			,					_lowercase))
								self.assertEqual(resulta[2]			,					1)

								A_       		=   send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}			,					_lowercase)
								self.assertIsInstance(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta['a'].cpu()			,					_lowercase))
								self.assertIsInstance(resulta['b']			,					_lowercase)
								self.assertTrue(torch.equal(resulta['b'][0].cpu()			,					_lowercase))
								self.assertTrue(torch.equal(resulta['b'][1].cpu()			,					_lowercase))
								self.assertEqual(resulta['c']			,					1)

								A_       		=   send_to_device(ExampleNamedTuple(a=_lowercase			,					b=[tensor, tensor]			,					c=1)			,					_lowercase)
								self.assertIsInstance(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta.a.cpu()			,					_lowercase))
								self.assertIsInstance(resulta.b			,					_lowercase)
								self.assertTrue(torch.equal(resulta.b[0].cpu()			,					_lowercase))
								self.assertTrue(torch.equal(resulta.b[1].cpu()			,					_lowercase))
								self.assertEqual(resulta.c			,					1)

								A_       		=   send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1})			,					_lowercase)
								self.assertIsInstance(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta['a'].cpu()			,					_lowercase))
								self.assertIsInstance(resulta['b']			,					_lowercase)
								self.assertTrue(torch.equal(resulta['b'][0].cpu()			,					_lowercase))
								self.assertTrue(torch.equal(resulta['b'][1].cpu()			,					_lowercase))
								self.assertEqual(resulta['c']			,					1)



					def 						__snake_case	(  self    :	Dict)						->       str:
								with self.assertRaises(_lowercase) as cm:
											A_       		=   recursively_apply(torch.tensor			,					(torch.tensor(1), 1)			,					error_on_other_type=_lowercase)
								self.assertEqual(
								    str(cm.exception)			,					'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'			,					)



					def 						__snake_case	(  self    :	Union[str, Any])						->       List[str]:
								A_       		=   torch.tensor([1, 2, 3, 4, 5])
								self.assertEqual(listify(_lowercase)			,					[1, 2, 3, 4, 5])

								A_       		=   torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
								self.assertEqual(listify(_lowercase)			,					[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])

								A_       		=   torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]])
								self.assertEqual(
								    listify(_lowercase)			,					[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]])



					def 						__snake_case	(  self    :	List[Any])						->       Union[str, Any]:
								with patch_environment(aa=1			,					BB=2):
											self.assertEqual(os.environ.get('AA')			,					'1')
											self.assertEqual(os.environ.get('BB')			,					'2')

								self.assertNotIn('AA'			,					os.environ)
								self.assertNotIn('BB'			,					os.environ)



					def 						__snake_case	(  self    :	Any)						->       str:
								A_       		=   RegressionModel()
								A_       		=   model.forward
								A_       		=   convert_outputs_to_fpaa(model.forward)
								A_       		=   extract_model_from_parallel(_lowercase			,					keep_fpaa_wrapper=_lowercase)
								A_       		=   pickle.dumps(_lowercase)



					@require_cuda
					def 						__snake_case	(  self    :	Union[str, Any])						->       Union[str, Any]:
								A_       		=   RegressionModel()
								A_       		=   model.forward
								A_       		=   torch.cuda.amp.autocast(dtype=torch.floataa)(model.forward)
								A_       		=   convert_outputs_to_fpaa(model.forward)
								A_       		=   extract_model_from_parallel(_lowercase			,					keep_fpaa_wrapper=_lowercase)
								A_       		=   pickle.dumps(_lowercase)



					@require_cuda
					@require_torch_min_version(version='2.0')
					def 						__snake_case	(  self    :	Optional[Any])						->       Dict:
								A_       		=   RegressionModel()
								A_       		=   model.forward
								A_       		=   torch.cuda.amp.autocast(dtype=torch.floataa)(model.forward)
								A_       		=   convert_outputs_to_fpaa(model.forward)
								A_       		=   torch.compile(model.forward			,					backend='inductor')
								A_       		=   torch.randn(4			,					10).cuda()
								A_       		=   model(_lowercase)



					def 						__snake_case	(  self    :	Union[str, Any])						->       Optional[Any]:
								A_       		=   RegressionModel()
								# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
								A_       		=   torch.nn.parallel.DataParallel(_lowercase)
								A_       		=   extract_model_from_parallel(_lowercase)

								self.assertEqual(_lowercase			,					_lowercase)



					@require_torch_min_version(version='2.0')
					def 						__snake_case	(  self    :	str)						->       int:
								A_       		=   RegressionModel()
								A_       		=   torch.compile(_lowercase)

								# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
								A_       		=   torch.nn.parallel.DataParallel(_lowercase)
								A_       		=   torch.compile(_lowercase)
								A_       		=   extract_model_from_parallel(_lowercase)

								self.assertEqual(compiled_model._orig_mod			,					compiled_model_unwrapped._orig_mod)



					def 						__snake_case	(  self    :	List[str])						->       Optional[int]:
								self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3])])			,					torch.device('cpu'))
								self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3])})			,					torch.device('cpu'))
								self.assertIsNone(find_device([1, 'a']))




