




import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment


def        _a  (     ) ->   Optional[Any]:
						"""simple docstring"""







						lowerCAmelCase__       =						torch.nn.Linear(2    ,   4   )
						lowerCAmelCase__       =						torch.optim.AdamW(model.parameters()    ,   lr=1.0   )
						lowerCAmelCase__       =						torch.optim.lr_scheduler.OneCycleLR(UpperCamelCase_    ,   max_lr=0.01    ,   steps_per_epoch=2    ,   epochs=1   )
						lowerCAmelCase__       =						DataLoader(TensorDataset(torch.tensor([1, 2, 3]   )   )   )
						lowerCAmelCase__       =						DataLoader(TensorDataset(torch.tensor([4, 5, 6]   )   )   )

						return model, optimizer, scheduler, train_dl, valid_dl


def        _a  (     UpperCamelCase_					:				Optional[Any]   ) ->   List[Any]:
						"""simple docstring"""







						return (model.weight.abs().sum() + model.bias.abs().sum()).item()


def        _a  (     UpperCamelCase_					:				str   ) ->   str:
						"""simple docstring"""







						lowerCAmelCase__       =						torch.nn.Linear(*tuple(model.weight.T.shape   )   ).state_dict()
						model.load_state_dict(UpperCamelCase_   )





class      lowercase__     (			_UpperCAmelCase	):

	@require_cuda
	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						Accelerator()
							assert PartialState._shared_state["_cpu"] is False
							assert PartialState._shared_state["device"].type == "cuda"
							with self.assertRaises(__UpperCAmelCase       ):
													lowerCAmelCase__       =						Accelerator(cpu=__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						Accelerator()
							lowerCAmelCase__       =						GradientState()
							assert state.num_steps == 1
							lowerCAmelCase__       =						4
							assert state.num_steps == 4

							assert state.sync_gradients is True
							lowerCAmelCase__       =						False
							assert state.sync_gradients is False
							GradientState._reset_state()

	def 		UpperCAmelCase  (							self       )->					Optional[int]:







							'''simple docstring'''





							lowerCAmelCase__       =						Accelerator()
							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						create_components()

							(
							    (
							    lowerCAmelCase__							
							)      ,							(
							    lowerCAmelCase__							
							)      ,							(
							    lowerCAmelCase__							
							)      ,							(
							    lowerCAmelCase__							
							)      ,							(
							    lowerCAmelCase__							
							)      ,														
							)       =						accelerator.prepare(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )

							self.assertTrue(prepared_model in accelerator._models       )
							self.assertTrue(prepared_optimizer in accelerator._optimizers       )
							self.assertTrue(prepared_scheduler in accelerator._schedulers       )
							self.assertTrue(prepared_train_dl in accelerator._dataloaders       )
							self.assertTrue(prepared_valid_dl in accelerator._dataloaders       )

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						Accelerator()
							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						create_components()
							accelerator.prepare(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )
							accelerator.free_memory()

							self.assertTrue(len(accelerator._models       ) == 0       )
							self.assertTrue(len(accelerator._optimizers       ) == 0       )
							self.assertTrue(len(accelerator._schedulers       ) == 0       )
							self.assertTrue(len(accelerator._dataloaders       ) == 0       )

	def 		UpperCAmelCase  (							self       )->					List[str]:







							'''simple docstring'''





							PartialState._reset_state()

							# Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
							def noop(*__UpperCAmelCase		,    **__UpperCAmelCase       ):
													pass

							with patch("torch.cuda.set_device"		,    __UpperCAmelCase       ), patch_environment(ACCELERATE_TORCH_DEVICE="cuda:64"       ):
													lowerCAmelCase__       =						Accelerator()
													self.assertEqual(str(accelerator.state.device       )		,    "cuda:64"       )

	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						Accelerator()
							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						create_components()
							accelerator.prepare(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )

							lowerCAmelCase__       =						get_signature(__UpperCAmelCase       )

							with tempfile.TemporaryDirectory() as tmpdirname:
													accelerator.save_state(__UpperCAmelCase       )

													# make sure random weights don't match
													load_random_weights(__UpperCAmelCase       )
													self.assertTrue(abs(model_signature - get_signature(__UpperCAmelCase       )       ) > 1E-3       )

													# make sure loaded weights match
													accelerator.load_state(__UpperCAmelCase       )
													self.assertTrue(abs(model_signature - get_signature(__UpperCAmelCase       )       ) < 1E-3       )

	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							lowerCAmelCase__       =						Accelerator()
							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						create_components()
							accelerator.prepare(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )

							lowerCAmelCase__       =						get_signature(__UpperCAmelCase       )

							# saving hook
							def save_config(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       ):
													lowerCAmelCase__       =						{"class_name": models[0].__class__.__name__}

													with open(os.path.join(__UpperCAmelCase		,    "data.json"       )		,    "w"       ) as f:
																			json.dump(__UpperCAmelCase		,    __UpperCAmelCase       )

        # loading hook
							def load_config(__UpperCAmelCase		,    __UpperCAmelCase       ):
													with open(os.path.join(__UpperCAmelCase		,    "data.json"       )		,    "r"       ) as f:
																			lowerCAmelCase__       =						json.load(__UpperCAmelCase       )

													lowerCAmelCase__       =						config["class_name"]

							lowerCAmelCase__       =						accelerator.register_save_state_pre_hook(__UpperCAmelCase       )
							lowerCAmelCase__       =						accelerator.register_load_state_pre_hook(__UpperCAmelCase       )

							with tempfile.TemporaryDirectory() as tmpdirname:
													accelerator.save_state(__UpperCAmelCase       )

													# make sure random weights don't match with hooks
													load_random_weights(__UpperCAmelCase       )
													self.assertTrue(abs(model_signature - get_signature(__UpperCAmelCase       )       ) > 1E-3       )

													# random class name to verify correct one is loaded
													lowerCAmelCase__       =						"random"

													# make sure loaded weights match with hooks
													accelerator.load_state(__UpperCAmelCase       )
													self.assertTrue(abs(model_signature - get_signature(__UpperCAmelCase       )       ) < 1E-3       )

													# mode.class_name is loaded from config
													self.assertTrue(model.class_name == model.__class__.__name__       )

							# remove hooks
							save_hook.remove()
							load_hook.remove()

							with tempfile.TemporaryDirectory() as tmpdirname:
													accelerator.save_state(__UpperCAmelCase       )

													# make sure random weights don't match with hooks removed
													load_random_weights(__UpperCAmelCase       )
													self.assertTrue(abs(model_signature - get_signature(__UpperCAmelCase       )       ) > 1E-3       )

													# random class name to verify correct one is loaded
													lowerCAmelCase__       =						"random"

													# make sure loaded weights match with hooks removed
													accelerator.load_state(__UpperCAmelCase       )
													self.assertTrue(abs(model_signature - get_signature(__UpperCAmelCase       )       ) < 1E-3       )

													# mode.class_name is NOT loaded from config
													self.assertTrue(model.class_name != model.__class__.__name__       )

	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						Accelerator()
							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						create_components()
							lowerCAmelCase__       =						None

							# This should work
							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						accelerator.prepare(
							    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )
							self.assertTrue(dummy_obj is None       )

	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						Accelerator()
							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						create_components()
							lowerCAmelCase__       =						[1, 2, 3]

							# This should work
							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__      ,							lowerCAmelCase__       =						accelerator.prepare(
							    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )
							self.assertEqual(
							    getattr(__UpperCAmelCase		,    "_is_accelerate_prepared"		,    __UpperCAmelCase       )		,    __UpperCAmelCase		,    "Dummy object should have `_is_accelerate_prepared` set to `True`"		,    )
							self.assertEqual(
							    getattr(__UpperCAmelCase		,    "_is_accelerate_prepared"		,    __UpperCAmelCase       )		,    __UpperCAmelCase		,    "Model is missing `_is_accelerator_prepared` or is set to `False`"		,    )
							self.assertEqual(
							    getattr(__UpperCAmelCase		,    "_is_accelerate_prepared"		,    __UpperCAmelCase       )		,    __UpperCAmelCase		,    "Optimizer is missing `_is_accelerator_prepared` or is set to `False`"		,    )
							self.assertEqual(
							    getattr(__UpperCAmelCase		,    "_is_accelerate_prepared"		,    __UpperCAmelCase       )		,    __UpperCAmelCase		,    "Scheduler is missing `_is_accelerator_prepared` or is set to `False`"		,    )
							self.assertEqual(
							    getattr(__UpperCAmelCase		,    "_is_accelerate_prepared"		,    __UpperCAmelCase       )		,    __UpperCAmelCase		,    "Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`"		,    )
							self.assertEqual(
							    getattr(__UpperCAmelCase		,    "_is_accelerate_prepared"		,    __UpperCAmelCase       )		,    __UpperCAmelCase		,    "Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`"		,    )

	@slow
	@require_bnb
	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							from transformers import AutoModelForCausalLM

							lowerCAmelCase__       =						AutoModelForCausalLM.from_pretrained(
							    "EleutherAI/gpt-neo-125m"		,    load_in_abit=__UpperCAmelCase		,    device_map={"": 0}		,    )
							lowerCAmelCase__       =						Accelerator()

							# This should work
							lowerCAmelCase__       =						accelerator.prepare(__UpperCAmelCase       )

	@slow
	@require_bnb
	def 		UpperCAmelCase  (							self       )->					Tuple:







							'''simple docstring'''





							from transformers import AutoModelForCausalLM

							lowerCAmelCase__       =						Accelerator()

							with init_empty_weights():
													lowerCAmelCase__       =						AutoModelForCausalLM.from_pretrained(
													    "EleutherAI/gpt-neo-125m"		,    )
													model.tie_weights()
													lowerCAmelCase__       =						infer_auto_device_map(__UpperCAmelCase       )
													lowerCAmelCase__       =						"cpu"

							lowerCAmelCase__       =						AutoModelForCausalLM.from_pretrained(
							    "EleutherAI/gpt-neo-125m"		,    device_map=__UpperCAmelCase		,    load_in_abit=__UpperCAmelCase		,    llm_inta_enable_fpaa_cpu_offload=__UpperCAmelCase       )

							# This should not work and get value error
							with self.assertRaises(__UpperCAmelCase       ):
													lowerCAmelCase__       =						accelerator.prepare(__UpperCAmelCase       )

	@slow
	@require_bnb
	@require_multi_gpu
	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							from transformers import AutoModelForCausalLM

							lowerCAmelCase__       =						{"distributed_type": DistributedType.MULTI_GPU}

							with init_empty_weights():
													lowerCAmelCase__       =						AutoModelForCausalLM.from_pretrained(
													    "EleutherAI/gpt-neo-125m"		,    )
													model.tie_weights()
													lowerCAmelCase__       =						infer_auto_device_map(__UpperCAmelCase       )
													lowerCAmelCase__       =						1

							lowerCAmelCase__       =						AutoModelForCausalLM.from_pretrained(
							    "EleutherAI/gpt-neo-125m"		,    load_in_abit=__UpperCAmelCase		,    device_map=__UpperCAmelCase		,    )
							lowerCAmelCase__       =						Accelerator()

							# This should not work and get value error
							with self.assertRaises(__UpperCAmelCase       ):
													lowerCAmelCase__       =						accelerator.prepare(__UpperCAmelCase       )

							PartialState._reset_state()

	@slow
	@require_bnb
	@require_multi_gpu
	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							from transformers import AutoModelForCausalLM

							with init_empty_weights():
													lowerCAmelCase__       =						AutoModelForCausalLM.from_pretrained(
													    "EleutherAI/gpt-neo-125m"		,    )
													lowerCAmelCase__       =						infer_auto_device_map(__UpperCAmelCase       )
													lowerCAmelCase__       =						1

							lowerCAmelCase__       =						AutoModelForCausalLM.from_pretrained(
							    "EleutherAI/gpt-neo-125m"		,    load_in_abit=__UpperCAmelCase		,    device_map=__UpperCAmelCase		,    )
							lowerCAmelCase__       =						Accelerator()

							# This should work
							lowerCAmelCase__       =						accelerator.prepare(__UpperCAmelCase       )


	@require_cuda
	def 		UpperCAmelCase  (							self       )->					List[str]:







							'''simple docstring'''





							lowerCAmelCase__       =						torch.nn.Linear(10		,    10       )
							lowerCAmelCase__       =						torch.optim.SGD(model.parameters()		,    lr=0.01       )
							lowerCAmelCase__       =						Accelerator(cpu=__UpperCAmelCase       )
							lowerCAmelCase__       =						accelerator.prepare(__UpperCAmelCase       )














def        _a  (     UpperCamelCase_					:				list    ,   UpperCamelCase_					:				list   ) ->   float:
						"""simple docstring"""







						_validate_point(UpperCamelCase_   )
						_validate_point(UpperCamelCase_   )
						if len(UpperCamelCase_   ) != len(UpperCamelCase_   ):
												raise ValueError("Both points must be in the same n-dimensional space"   )

						return float(sum(abs(a - b   ) for a, b in zip(UpperCamelCase_    ,   UpperCamelCase_   )   )   )


def        _a  (     UpperCamelCase_					:				list[float]   ) ->   None:
						"""simple docstring"""







						if point:
												if isinstance(UpperCamelCase_    ,   UpperCamelCase_   ):
																		for item in point:
																								if not isinstance(UpperCamelCase_    ,   (int, float)   ):
																														lowerCAmelCase__       =						(
																														    "Expected a list of numbers as input, found "
																														    F"{type(UpperCamelCase_   ).__name__}"
																														)
																														raise TypeError(UpperCamelCase_   )
												else:
																		lowerCAmelCase__       =						F"Expected a list of numbers as input, found {type(UpperCamelCase_   ).__name__}"
																		raise TypeError(UpperCamelCase_   )
						else:
												raise ValueError("Missing an input"   )


def        _a  (     UpperCamelCase_					:				list    ,   UpperCamelCase_					:				list   ) ->   float:
						"""simple docstring"""







						_validate_point(UpperCamelCase_   )
						_validate_point(UpperCamelCase_   )
						if len(UpperCamelCase_   ) != len(UpperCamelCase_   ):
												raise ValueError("Both points must be in the same n-dimensional space"   )

						return float(sum(abs(x - y   ) for x, y in zip(UpperCamelCase_    ,   UpperCamelCase_   )   )   )


if __name__ == "__main__":
	import doctest

	doctest.testmod()






