

'''simple docstring'''

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






     '''simple docstring'''





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

     return model, optimizer, scheduler, train_dl, valid_dl

def 					__UpperCamelCase				(						lowercase__    :	Optional[int]			):






     '''simple docstring'''





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

def 					__UpperCamelCase				(						lowercase__    :	Tuple			):






     '''simple docstring'''





     __lowercase      		=torch.nn.Linear(*tuple(model.weight.T.shape			)			).state_dict()
     model.load_state_dict(lowercase__			)







class 	lowerCAmelCase  (   A ):







       @require_cuda
       def       snake_case (				self      :			int   ):







            """simple docstring"""







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







       def       snake_case (				self      :			List[Any]   ):







            """simple docstring"""







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

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







       def       snake_case (				self      :			List[str]   ):







            """simple docstring"""







            __lowercase      		=Accelerator()
            __lowercase		,  __lowercase		,  __lowercase		,  __lowercase		,  __lowercase      		=create_components()

            (
                (
                __lowercase
            )		,  (
                __lowercase
            )		,  (
                __lowercase
            )		,  (
                __lowercase
            )		,  (
                __lowercase
            )		,  
            )      		=accelerator.prepare(__lowercase    ,   __lowercase    ,   __lowercase    ,   __lowercase    ,   __lowercase   )

            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       snake_case (				self      :			Tuple   ):







            """simple docstring"""







            __lowercase      		=Accelerator()
            __lowercase		,  __lowercase		,  __lowercase		,  __lowercase		,  __lowercase      		=create_components()
            accelerator.prepare(__lowercase    ,   __lowercase    ,   __lowercase    ,   __lowercase    ,   __lowercase   )
            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       snake_case (				self      :			Dict   ):







            """simple docstring"""







            PartialState._reset_state()

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

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







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







            """simple docstring"""







            __lowercase      		=Accelerator()
            __lowercase		,  __lowercase		,  __lowercase		,  __lowercase		,  __lowercase      		=create_components()
            accelerator.prepare(__lowercase    ,   __lowercase    ,   __lowercase    ,   __lowercase    ,   __lowercase   )

            __lowercase      		=get_signature(__lowercase   )

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

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

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







       def       snake_case (				self      :			Optional[int]   ):







            """simple docstring"""







            __lowercase      		=Accelerator()
            __lowercase		,  __lowercase		,  __lowercase		,  __lowercase		,  __lowercase      		=create_components()
            accelerator.prepare(__lowercase    ,   __lowercase    ,   __lowercase    ,   __lowercase    ,   __lowercase   )

            __lowercase      		=get_signature(__lowercase   )

            # saving hook
            def save_config(__lowercase      :			List[Any]    ,   __lowercase      :			str    ,   __lowercase      :			str   ):
                 __lowercase      		={'class_name': models[0].__class__.__name__}

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

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

                 __lowercase      		=config['class_name']

            __lowercase      		=accelerator.register_save_state_pre_hook(__lowercase   )
            __lowercase      		=accelerator.register_load_state_pre_hook(__lowercase   )

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

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

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

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

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

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

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

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







       def       snake_case (				self      :			List[Any]   ):







            """simple docstring"""







            __lowercase      		=Accelerator()
            __lowercase		,  __lowercase		,  __lowercase		,  __lowercase		,  __lowercase      		=create_components()
            __lowercase      		=None

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







       def       snake_case (				self      :			List[str]   ):







            """simple docstring"""







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

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







       @slow
       @require_bnb
       def       snake_case (				self      :			int   ):







            """simple docstring"""







            from transformers import AutoModelForCausalLM

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

            # This should work
            __lowercase      		=accelerator.prepare(__lowercase   )







       @slow
       @require_bnb
       def       snake_case (				self      :			Union[str, Any]   ):







            """simple docstring"""







            from transformers import AutoModelForCausalLM

            __lowercase      		=Accelerator()

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

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

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







       @slow
       @require_bnb
       @require_multi_gpu
       def       snake_case (				self      :			Tuple   ):







            """simple docstring"""







            from transformers import AutoModelForCausalLM

            __lowercase      		={'distributed_type': DistributedType.MULTI_GPU}

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

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

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

            PartialState._reset_state()







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







            """simple docstring"""







            from transformers import AutoModelForCausalLM

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

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

            # This should work
            __lowercase      		=accelerator.prepare(__lowercase   )







       @require_cuda
       def       snake_case (				self      :			Any   ):







            """simple docstring"""







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





'''simple docstring'''

# Function to print upper half of diamond (pyramid)

def 					__UpperCamelCase				(						lowercase__    :	Optional[Any]			):






     '''simple docstring'''





     for i in range(0,       lowercase__			):
          for _ in range(0,       n - i - 1			):  # printing spaces
               print(' ',       end=''			)
          for _ in range(0,       i + 1			):  # printing stars
               print('* ',       end=''			)
          print()

def 					__UpperCamelCase				(						lowercase__    :	Optional[Any]			):






     '''simple docstring'''





     for i in range(lowercase__,       0,       -1			):
          for _ in range(lowercase__,       0,       -1			):  # printing stars
               print('* ',       end=''			)
          print()
          for _ in range(n - i + 1,       0,       -1			):  # printing spaces
               print(' ',       end=''			)

def 					__UpperCamelCase				(						lowercase__    :	Any			):






     '''simple docstring'''





     if n <= 0:
          print('       ...       ....        nothing printing :('			)
          return
     floyd(lowercase__			)  # upper half
     reverse_floyd(lowercase__			)  # lower half


if __name__ == "__main__":
       print(r'''| /\ | |- |  |-  |--| |\  /| |-''')
       print(r'''|/  \| |- |_ |_  |__| | \/ | |_''')
       UpperCAmelCase					 =	1
       while K:
              UpperCAmelCase					 =	int(input('''enter the number and , and see the magic : '''))
              print()
              pretty_print(user_number)
              UpperCAmelCase					 =	int(input('''press 0 to exit... and 1 to continue...'''))

       print('''Good Bye...''')


