






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

 return model, optimizer, scheduler, train_dl, valid_dl
def 		_a       ( a					:Dict					)				->     Tuple:
 return (model.weight.abs().sum() + model.bias.abs().sum()).item()
def 		_a       ( a					:int					)				->     int:
 a						    = torch.nn.Linear(*tuple(model.weight.T.shape					)					).state_dict()
 model.load_state_dict(a					)






class 						lowercase_   (      lowercase  ):






      '''simple docstring'''




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




       """simple docstring"""


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







      def  __lowerCAmelCase   ( self :		Dict    )       ->Dict:




       """simple docstring"""


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

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







      def  __lowerCAmelCase   ( self :		Dict    )       ->Any:




       """simple docstring"""


       a						    = Accelerator()
       a    ,	a    ,	a    ,	a    ,	a						    = create_components()

       (
           (
           a
       )    ,	(
           a
       )    ,	(
           a
       )    ,	(
           a
       )    ,	(
           a
       )    ,	
       )						    = 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  __lowerCAmelCase   ( self :		Tuple    )       ->str:




       """simple docstring"""


       a						    = Accelerator()
       a    ,	a    ,	a    ,	a    ,	a						    = 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  __lowerCAmelCase   ( self :		Optional[int]    )       ->Union[str, Any]:




       """simple docstring"""


       PartialState._reset_state()

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

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







      def  __lowerCAmelCase   ( self :		Dict    )       ->Dict:




       """simple docstring"""


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

       a						    = 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  __lowerCAmelCase   ( self :		int    )       ->Any:




       """simple docstring"""


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

       a						    = get_signature(__UpperCAmelCase    )

       # saving hook
       def save_config(__UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Tuple					,   __UpperCAmelCase :		List[Any]    ):
        a						    = {'''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 :		str					,   __UpperCAmelCase :		str    ):
        with open(os.path.join(__UpperCAmelCase					,   '''data.json'''    )					,   '''r'''    ) as f:
         a						    = json.load(__UpperCAmelCase    )

        a						    = config['''class_name''']

       a						    = accelerator.register_save_state_pre_hook(__UpperCAmelCase    )
       a						    = 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
        a						    = '''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
        a						    = '''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  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->Dict:




       """simple docstring"""


       a						    = Accelerator()
       a    ,	a    ,	a    ,	a    ,	a						    = create_components()
       a						    = None

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







      def  __lowerCAmelCase   ( self :		Dict    )       ->List[Any]:




       """simple docstring"""


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

       # This should work
       a    ,	a    ,	a    ,	a    ,	a    ,	a						    = 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  __lowerCAmelCase   ( self :		Any    )       ->Tuple:




       """simple docstring"""


       from transformers import AutoModelForCausalLM

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

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







      @slow
      @require_bnb
      def  __lowerCAmelCase   ( self :		str    )       ->Tuple:




       """simple docstring"""


       from transformers import AutoModelForCausalLM

       a						    = Accelerator()

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

       a						    = 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    ):
        a						    = accelerator.prepare(__UpperCAmelCase    )







      @slow
      @require_bnb
      @require_multi_gpu
      def  __lowerCAmelCase   ( self :		int    )       ->List[str]:




       """simple docstring"""


       from transformers import AutoModelForCausalLM

       a						    = {'''distributed_type''': DistributedType.MULTI_GPU}

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

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

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

       PartialState._reset_state()







      @slow
      @require_bnb
      @require_multi_gpu
      def  __lowerCAmelCase   ( self :		int    )       ->Optional[int]:




       """simple docstring"""


       from transformers import AutoModelForCausalLM

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

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

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







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




       """simple docstring"""


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










from typing import Any

import numpy as np
def 		_a       ( a					:np.ndarray					)				->     bool:
 return np.array_equal(a      ,				matrix.conjugate().T					)
def 		_a       ( a					:np.ndarray      ,				a					:np.ndarray					)				->     Any:
 a						    = v.conjugate().T
 a						    = v_star.dot(a					)
 assert isinstance(a      ,				np.ndarray					)
 return (v_star_dot.dot(a					)) / (v_star.dot(a					))
def 		_a       ( )				->     None:
 a						    = np.array([[2, 2 + 1j, 4], [2 - 1j, 3, 1j], [4, -1j, 1]]					)
 a						    = np.array([[1], [2], [3]]					)
 assert is_hermitian(a					), F"""{a} is not hermitian."""
 print(rayleigh_quotient(a      ,				a					)					)

 a						    = np.array([[1, 2, 4], [2, 3, -1], [4, -1, 1]]					)
 assert is_hermitian(a					), F"""{a} is not hermitian."""
 assert rayleigh_quotient(a      ,				a					) == float(3					)


if __name__ == "__main__":
     import doctest

     doctest.testmod()
     tests()


