


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 				UpperCamelCase_(      )    ->						Dict:
    _lowercase       :							List[Any]      		=		torch.nn.Linear(2							, 4 )
    _lowercase       :							Optional[int]      		=		torch.optim.AdamW(model.parameters()							, lr=1.0 )
    _lowercase       :							int      		=		torch.optim.lr_scheduler.OneCycleLR(lowerCamelCase_							, max_lr=0.01							, steps_per_epoch=2							, epochs=1 )
    _lowercase       :							List[str]      		=		DataLoader(TensorDataset(torch.tensor([1, 2, 3] ) ) )
    _lowercase       :							Tuple      		=		DataLoader(TensorDataset(torch.tensor([4, 5, 6] ) ) )

    return model, optimizer, scheduler, train_dl, valid_dl



def 				UpperCamelCase_(      lowerCamelCase_ )    ->						List[str]:
    return (model.weight.abs().sum() + model.bias.abs().sum()).item()



def 				UpperCamelCase_(      lowerCamelCase_ )    ->						List[Any]:
    _lowercase       :							Tuple      		=		torch.nn.Linear(*tuple(model.weight.T.shape ) ).state_dict()
    model.load_state_dict(lowerCamelCase_ )






class _lowerCamelCase(						UpperCAmelCase__   ):
       @require_cuda
       def 			UpperCamelCase     (					self)       ->   Tuple:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		Accelerator()
           assert PartialState._shared_state["_cpu"] is False
           assert PartialState._shared_state["device"].type == "cuda"
           with self.assertRaises(_SCREAMING_SNAKE_CASE):
               _lowercase       :							Optional[int]      		=		Accelerator(cpu=_SCREAMING_SNAKE_CASE)
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase       :							Any      		=		Accelerator()
           _lowercase       :							List[str]      		=		GradientState()
           assert state.num_steps == 1
           _lowercase       :							Optional[Any]      		=		4
           assert state.num_steps == 4

           assert state.sync_gradients is True
           _lowercase       :							List[str]      		=		False
           assert state.sync_gradients is False
           GradientState._reset_state()
       def 			UpperCamelCase     (					self)       ->   Union[str, Any]:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		Accelerator()
           _lowercase       :							List[str]      		=		create_components()

           (
               _lowercase		
           )       :							Union[str, Any]      		=		accelerator.prepare(_SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE)

           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)       ->   Tuple:




           """simple docstring"""

           _lowercase       :							int      		=		Accelerator()
           _lowercase       :							Dict      		=		create_components()
           accelerator.prepare(_SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE)
           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[Any]:




           """simple docstring"""

           PartialState._reset_state()

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

           with patch('torch.cuda.set_device',      _SCREAMING_SNAKE_CASE), patch_environment(ACCELERATE_TORCH_DEVICE='cuda:64'):
               _lowercase       :							List[str]      		=		Accelerator()
               self.assertEqual(str(accelerator.state.device),      'cuda:64')
       def 			UpperCamelCase     (					self)       ->   str:




           """simple docstring"""

           _lowercase       :							Any      		=		Accelerator()
           _lowercase       :							int      		=		create_components()
           accelerator.prepare(_SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE)

           _lowercase       :							Any      		=		get_signature(_SCREAMING_SNAKE_CASE)

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

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

               # make sure loaded weights match
               accelerator.load_state(_SCREAMING_SNAKE_CASE)
               self.assertTrue(abs(model_signature - get_signature(_SCREAMING_SNAKE_CASE)) < 1E-3)
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		Accelerator()
           _lowercase       :							int      		=		create_components()
           accelerator.prepare(_SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE)

           _lowercase       :							Tuple      		=		get_signature(_SCREAMING_SNAKE_CASE)

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

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

                     # loading hook
           def load_config(lowerCamelCase,      lowerCamelCase):
               with open(os.path.join(_SCREAMING_SNAKE_CASE,      'data.json'),      'r') as f:
                   _lowercase       :							Union[str, Any]      		=		json.load(_SCREAMING_SNAKE_CASE)

               _lowercase       :							Union[str, Any]      		=		config["""class_name"""]

           _lowercase       :							Union[str, Any]      		=		accelerator.register_save_state_pre_hook(_SCREAMING_SNAKE_CASE)
           _lowercase       :							List[str]      		=		accelerator.register_load_state_pre_hook(_SCREAMING_SNAKE_CASE)

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

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

               # random class name to verify correct one is loaded
               _lowercase       :							Union[str, Any]      		=		"""random"""

               # make sure loaded weights match with hooks
               accelerator.load_state(_SCREAMING_SNAKE_CASE)
               self.assertTrue(abs(model_signature - get_signature(_SCREAMING_SNAKE_CASE)) < 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(_SCREAMING_SNAKE_CASE)

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

               # random class name to verify correct one is loaded
               _lowercase       :							List[Any]      		=		"""random"""

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

               # mode.class_name is NOT loaded from config
               self.assertTrue(model.class_name != model.__class__.__name__)
       def 			UpperCamelCase     (					self)       ->   str:




           """simple docstring"""

           _lowercase       :							List[str]      		=		Accelerator()
           _lowercase       :							Optional[int]      		=		create_components()
           _lowercase       :							List[str]      		=		None

           # This should work
           _lowercase       :							Optional[Any]      		=		accelerator.prepare(
               _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE)
           self.assertTrue(dummy_obj is None)
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           _lowercase       :							List[str]      		=		Accelerator()
           _lowercase       :							List[Any]      		=		create_components()
           _lowercase       :							Union[str, Any]      		=		[1, 2, 3]

           # This should work
           _lowercase       :							Optional[Any]      		=		accelerator.prepare(
               _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE,      _SCREAMING_SNAKE_CASE)
           self.assertEqual(
               getattr(_SCREAMING_SNAKE_CASE,      '_is_accelerate_prepared',      _SCREAMING_SNAKE_CASE),      _SCREAMING_SNAKE_CASE,      'Dummy object should have `_is_accelerate_prepared` set to `True`',      )
           self.assertEqual(
               getattr(_SCREAMING_SNAKE_CASE,      '_is_accelerate_prepared',      _SCREAMING_SNAKE_CASE),      _SCREAMING_SNAKE_CASE,      'Model is missing `_is_accelerator_prepared` or is set to `False`',      )
           self.assertEqual(
               getattr(_SCREAMING_SNAKE_CASE,      '_is_accelerate_prepared',      _SCREAMING_SNAKE_CASE),      _SCREAMING_SNAKE_CASE,      'Optimizer is missing `_is_accelerator_prepared` or is set to `False`',      )
           self.assertEqual(
               getattr(_SCREAMING_SNAKE_CASE,      '_is_accelerate_prepared',      _SCREAMING_SNAKE_CASE),      _SCREAMING_SNAKE_CASE,      'Scheduler is missing `_is_accelerator_prepared` or is set to `False`',      )
           self.assertEqual(
               getattr(_SCREAMING_SNAKE_CASE,      '_is_accelerate_prepared',      _SCREAMING_SNAKE_CASE),      _SCREAMING_SNAKE_CASE,      'Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`',      )
           self.assertEqual(
               getattr(_SCREAMING_SNAKE_CASE,      '_is_accelerate_prepared',      _SCREAMING_SNAKE_CASE),      _SCREAMING_SNAKE_CASE,      'Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`',      )
       @slow
       @require_bnb
       def 			UpperCamelCase     (					self)       ->   List[Any]:




           """simple docstring"""

           from transformers import AutoModelForCausalLM

           _lowercase       :							Any      		=		AutoModelForCausalLM.from_pretrained(
               'EleutherAI/gpt-neo-125m',      load_in_abit=_SCREAMING_SNAKE_CASE,      device_map={'': 0},      )
           _lowercase       :							Union[str, Any]      		=		Accelerator()

           # This should work
           _lowercase       :							str      		=		accelerator.prepare(_SCREAMING_SNAKE_CASE)
       @slow
       @require_bnb
       def 			UpperCamelCase     (					self)       ->   Tuple:




           """simple docstring"""

           from transformers import AutoModelForCausalLM

           _lowercase       :							Any      		=		Accelerator()

           with init_empty_weights():
               _lowercase       :							Optional[int]      		=		AutoModelForCausalLM.from_pretrained(
                   'EleutherAI/gpt-neo-125m',      )
               model.tie_weights()
               _lowercase       :							Any      		=		infer_auto_device_map(_SCREAMING_SNAKE_CASE)
               _lowercase       :							Optional[int]      		=		"""cpu"""

           _lowercase       :							List[str]      		=		AutoModelForCausalLM.from_pretrained(
               'EleutherAI/gpt-neo-125m',      device_map=_SCREAMING_SNAKE_CASE,      load_in_abit=_SCREAMING_SNAKE_CASE,      llm_inta_enable_fpaa_cpu_offload=_SCREAMING_SNAKE_CASE)

           # This should not work and get value error
           with self.assertRaises(_SCREAMING_SNAKE_CASE):
               _lowercase       :							Union[str, Any]      		=		accelerator.prepare(_SCREAMING_SNAKE_CASE)
       @slow
       @require_bnb
       @require_multi_gpu
       def 			UpperCamelCase     (					self)       ->   Optional[Any]:




           """simple docstring"""

           from transformers import AutoModelForCausalLM

           _lowercase       :							Dict      		=		{"""distributed_type""": DistributedType.MULTI_GPU}

           with init_empty_weights():
               _lowercase       :							Tuple      		=		AutoModelForCausalLM.from_pretrained(
                   'EleutherAI/gpt-neo-125m',      )
               model.tie_weights()
               _lowercase       :							List[str]      		=		infer_auto_device_map(_SCREAMING_SNAKE_CASE)
               _lowercase       :							List[Any]      		=		1

           _lowercase       :							Optional[int]      		=		AutoModelForCausalLM.from_pretrained(
               'EleutherAI/gpt-neo-125m',      load_in_abit=_SCREAMING_SNAKE_CASE,      device_map=_SCREAMING_SNAKE_CASE,      )
           _lowercase       :							Optional[int]      		=		Accelerator()

           # This should not work and get value error
           with self.assertRaises(_SCREAMING_SNAKE_CASE):
               _lowercase       :							Union[str, Any]      		=		accelerator.prepare(_SCREAMING_SNAKE_CASE)

           PartialState._reset_state()
       @slow
       @require_bnb
       @require_multi_gpu
       def 			UpperCamelCase     (					self)       ->   Optional[Any]:




           """simple docstring"""

           from transformers import AutoModelForCausalLM

           with init_empty_weights():
               _lowercase       :							List[str]      		=		AutoModelForCausalLM.from_pretrained(
                   'EleutherAI/gpt-neo-125m',      )
               _lowercase       :							str      		=		infer_auto_device_map(_SCREAMING_SNAKE_CASE)
               _lowercase       :							int      		=		1

           _lowercase       :							int      		=		AutoModelForCausalLM.from_pretrained(
               'EleutherAI/gpt-neo-125m',      load_in_abit=_SCREAMING_SNAKE_CASE,      device_map=_SCREAMING_SNAKE_CASE,      )
           _lowercase       :							int      		=		Accelerator()

           # This should work
           _lowercase       :							List[str]      		=		accelerator.prepare(_SCREAMING_SNAKE_CASE)




       @require_cuda
       def 			UpperCamelCase     (					self)       ->   Tuple:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		torch.nn.Linear(10,      10)
           _lowercase       :							int      		=		torch.optim.SGD(model.parameters(),      lr=0.0_1)
           _lowercase       :							str      		=		Accelerator(cpu=_SCREAMING_SNAKE_CASE)
           _lowercase       :							int      		=		accelerator.prepare(_SCREAMING_SNAKE_CASE)







"""simple docstring"""
def       _snake_case						(						UpperCamelCase						:		list						):
 if not isinstance(UpperCamelCase			,		UpperCamelCase						):
  raise ValueError("""Input series is not valid, valid series - [2, 4, 6]"""						)
 if len(UpperCamelCase						) == 0:
  raise ValueError("""Input list must be a non empty list"""						)
 if len(UpperCamelCase						) == 1:
  return True
 UpperCAmelCase	:					Union[str, Any]	=	series[1] - series[0]
 for index in range(len(UpperCamelCase						) - 1						):
  if series[index + 1] - series[index] != common_diff:
   return False
 return True

def       _snake_case						(						UpperCamelCase						:		list						):
 if not isinstance(UpperCamelCase			,		UpperCamelCase						):
  raise ValueError("""Input series is not valid, valid series - [2, 4, 6]"""						)
 if len(UpperCamelCase						) == 0:
  raise ValueError("""Input list must be a non empty list"""						)
 UpperCAmelCase	:					Any	=	0
 for val in series:
  answer += val
 return answer / len(UpperCamelCase						)


if __name__ == "__main__":
      import doctest

      doctest.testmod()
