





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

   return model, optimizer, scheduler, train_dl, valid_dl






def lowerCamelCase__    (							snake_case_	:  Any       )      ->				Any:
   return (model.weight.abs().sum() + model.bias.abs().sum()).item()






def lowerCamelCase__    (							snake_case_	:  Union[str, Any]       )      ->				Union[str, Any]:
   __snake_case							       =   torch.nn.Linear(*tuple(model.weight.T.shape       )       ).state_dict()
   model.load_state_dict(snake_case_       )







class 		SCREAMING_SNAKE_CASE__				(       _UpperCAmelCase  ):


      @require_cuda
      def    a (self       :      Tuple      ):





         """simple docstring"""



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


      def    a (self       :      Tuple      ):





         """simple docstring"""



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

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


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()

         (
             (
             __snake_case
         )       ,	(
             __snake_case
         )       ,	(
             __snake_case
         )       ,	(
             __snake_case
         )       ,	(
             __snake_case
         )       ,	
         )							       =   accelerator.prepare(a__				,			a__				,			a__				,			a__				,			a__      )

         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    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         accelerator.prepare(a__				,			a__				,			a__				,			a__				,			a__      )
         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    a (self       :      str      ):





         """simple docstring"""



         PartialState._reset_state()

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

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


      def    a (self       :      Union[str, Any]      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         accelerator.prepare(a__				,			a__				,			a__				,			a__				,			a__      )

         __snake_case							       =   get_signature(a__      )

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

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

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


      def    a (self       :      int      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         accelerator.prepare(a__				,			a__				,			a__				,			a__				,			a__      )

         __snake_case							       =   get_signature(a__      )

         # saving hook
         def save_config(a__       :      Any				,			a__       :      Any				,			a__       :      Optional[int]      ):
            __snake_case							       =   {'''class_name''': models[0].__class__.__name__}

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

        # loading hook
         def load_config(a__       :      str				,			a__       :      Optional[Any]      ):
            with open(os.path.join(a__				,			'''data.json'''      )				,			'''r'''      ) as f:
               __snake_case							       =   json.load(a__      )

            __snake_case							       =   config['''class_name''']

         __snake_case							       =   accelerator.register_save_state_pre_hook(a__      )
         __snake_case							       =   accelerator.register_load_state_pre_hook(a__      )

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

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

            # random class name to verify correct one is loaded
            __snake_case							       =   '''random'''

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

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

            # random class name to verify correct one is loaded
            __snake_case							       =   '''random'''

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

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


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         __snake_case							       =   None

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


      def    a (self       :      Union[str, Any]      ):





         """simple docstring"""



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

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


      @slow
      @require_bnb
      def    a (self       :      Optional[Any]      ):





         """simple docstring"""



         from transformers import AutoModelForCausalLM

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

         # This should work
         __snake_case							       =   accelerator.prepare(a__      )


      @slow
      @require_bnb
      def    a (self       :      Any      ):





         """simple docstring"""



         from transformers import AutoModelForCausalLM

         __snake_case							       =   Accelerator()

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

         __snake_case							       =   AutoModelForCausalLM.from_pretrained(
             '''EleutherAI/gpt-neo-125m'''				,			device_map=a__				,			load_in_abit=a__				,			llm_inta_enable_fpaa_cpu_offload=a__      )

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


      @slow
      @require_bnb
      @require_multi_gpu
      def    a (self       :      int      ):





         """simple docstring"""



         from transformers import AutoModelForCausalLM

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

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

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

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

         PartialState._reset_state()


      @slow
      @require_bnb
      @require_multi_gpu
      def    a (self       :      int      ):





         """simple docstring"""



         from transformers import AutoModelForCausalLM

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

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

         # This should work
         __snake_case							       =   accelerator.prepare(a__      )


      @require_cuda
      def    a (self       :      str      ):





         """simple docstring"""



         __snake_case							       =   torch.nn.Linear(10				,			10      )
         __snake_case							       =   torch.optim.SGD(model.parameters()				,			lr=0.0_1      )
         __snake_case							       =   Accelerator(cpu=a__      )
         __snake_case							       =   accelerator.prepare(a__      )








def lowerCamelCase__    (							snake_case_	:  int       )      ->				int:
   if not isinstance(snake_case_       ,      snake_case_       ) or number < 0:
      raise ValueError('''Input must be a non-negative integer'''       )

   __snake_case							       =   0
   while number:
      # This way we arrive at next set bit (next 1) instead of looping
      # through each bit and checking for 1s hence the
      # loop won't run 32 times it will only run the number of `1` times
      number &= number - 1
      count += 1
   return count


if __name__ == "__main__":
    import doctest

    doctest.testmod()

