from timeit import timeit

__lowercase										= {
    '''MALAYALAM''': True,
    '''String''': False,
    '''rotor''': True,
    '''level''': True,
    '''A''': True,
    '''BB''': True,
    '''ABC''': False,
    '''amanaplanacanalpanama''': True,  # "a man a plan a canal panama"
}
# Ensure our test data is valid
assert all((key == key[::-1]) is value for key, value in test_data.items())



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    __UpperCamelCase					:List[Any]        =		0
    __UpperCamelCase					:str        =		len(SCREAMING_SNAKE_CASE     ) - 1
    while start_i < end_i:
        if s[start_i] == s[end_i]:
            start_i += 1
            end_i -= 1
        else:
            return False
    return True



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    __UpperCamelCase					:List[Any]        =		len(SCREAMING_SNAKE_CASE     ) // 2
    __UpperCamelCase					:Dict        =		len(SCREAMING_SNAKE_CASE     )

    # We need to traverse till half of the length of string
    # as we can get access of the i'th last element from
    # i'th index.
    # eg: [0,1,2,3,4,5] => 4th index can be accessed
    # with the help of 1st index (i==n-i-1)
    # where n is length of string
    return all(s[i] == s[n - i - 1] for i in range(SCREAMING_SNAKE_CASE     )     )



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    if len(SCREAMING_SNAKE_CASE     ) <= 2:
        return True
    if s[0] == s[len(SCREAMING_SNAKE_CASE     ) - 1]:
        return is_palindrome_recursive(s[1:-1]     )
    else:
        return False



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    return s == s[::-1]



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    __UpperCamelCase					:Tuple        =		f"""all({name}(key) is value for key, value in test_data.items())"""
    __UpperCamelCase					:str        =		f"""from __main__ import test_data, {name}"""
    __UpperCamelCase					:List[Any]        =		500_000
    __UpperCamelCase					:List[str]        =		timeit(stmt=SCREAMING_SNAKE_CASE  ,      setup=SCREAMING_SNAKE_CASE  ,      number=SCREAMING_SNAKE_CASE     )
    print(f"""{name:<35} finished {number:,} runs in {result:.5f} seconds"""     )


if __name__ == "__main__":
    for key, value in test_data.items():
        assert is_palindrome(key) is is_palindrome_recursive(key)
        assert is_palindrome(key) is is_palindrome_slice(key)
        print(F'{key:21} {value}')
    print('''a man a plan a canal panama''')

    # finished 500,000 runs in 0.46793 seconds
    benchmark_function('''is_palindrome_slice''')
    # finished 500,000 runs in 0.85234 seconds
    benchmark_function('''is_palindrome''')
    # finished 500,000 runs in 1.32028 seconds
    benchmark_function('''is_palindrome_recursive''')
    # finished 500,000 runs in 2.08679 seconds
    benchmark_function('''is_palindrome_traversal''')






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  (       ):
    '''simple docstring'''







    __UpperCamelCase					:List[str]        =		torch.nn.Linear(2  ,      4     )
    __UpperCamelCase					:Any        =		torch.optim.AdamW(model.parameters()  ,      lr=1.0     )
    __UpperCamelCase					:List[Any]        =		torch.optim.lr_scheduler.OneCycleLR(SCREAMING_SNAKE_CASE  ,      max_lr=0.01  ,      steps_per_epoch=2  ,      epochs=1     )
    __UpperCamelCase					:List[Any]        =		DataLoader(TensorDataset(torch.tensor([1, 2, 3]     )     )     )
    __UpperCamelCase					:Dict        =		DataLoader(TensorDataset(torch.tensor([4, 5, 6]     )     )     )

    return model, optimizer, scheduler, train_dl, valid_dl



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







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



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    __UpperCamelCase					:Union[str, Any]        =		torch.nn.Linear(*tuple(model.weight.T.shape     )     ).state_dict()
    model.load_state_dict(SCREAMING_SNAKE_CASE     )


class      lowerCamelCase_					(       UpperCAmelCase_			):






    '''simple docstring'''






    @require_cuda
    def       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:Dict        =		Accelerator()
        assert PartialState._shared_state["_cpu"] is False
        assert PartialState._shared_state["device"].type == "cuda"
        with self.assertRaises(__lowercase):
            __UpperCamelCase					:Any        =		Accelerator(cpu=__lowercase)







    def       UpperCamelCase__       (     self)	->     Any:
        __UpperCamelCase					:List[Any]        =		Accelerator()
        __UpperCamelCase					:List[Any]        =		GradientState()
        assert state.num_steps == 1
        __UpperCamelCase					:Any        =		4
        assert state.num_steps == 4

        assert state.sync_gradients is True
        __UpperCamelCase					:int        =		False
        assert state.sync_gradients is False
        GradientState._reset_state()







    def       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:Tuple        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Optional[Any]        =		create_components()

        (
            (
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	
        )					:int        =		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       UpperCamelCase__       (     self)	->     Optional[Any]:
        __UpperCamelCase					:str        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Optional[Any]        =		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       UpperCamelCase__       (     self)	->     Union[str, Any]:
        PartialState._reset_state()

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

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







    def       UpperCamelCase__       (     self)	->     List[Any]:
        __UpperCamelCase					:List[Any]        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:List[Any]        =		create_components()
        accelerator.prepare(__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)

        __UpperCamelCase					:Tuple        =		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       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:List[Any]        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Any        =		create_components()
        accelerator.prepare(__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)

        __UpperCamelCase					:Any        =		get_signature(__lowercase)

        # saving hook
        def save_config(__lowercase			,			__lowercase			,			__lowercase):
            __UpperCamelCase					:Union[str, Any]        =		{'''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			,			__lowercase):
            with open(os.path.join(__lowercase			,			'''data.json''')			,			'''r''') as f:
                __UpperCamelCase					:Dict        =		json.load(__lowercase)

            __UpperCamelCase					:Dict        =		config['''class_name''']

        __UpperCamelCase					:Union[str, Any]        =		accelerator.register_save_state_pre_hook(__lowercase)
        __UpperCamelCase					:Any        =		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
            __UpperCamelCase					:int        =		'''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
            __UpperCamelCase					:Dict        =		'''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       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:Optional[Any]        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Union[str, Any]        =		create_components()
        __UpperCamelCase					:Optional[Any]        =		None

        # This should work
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:List[Any]        =		accelerator.prepare(
            __lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)
        self.assertTrue(dummy_obj is None)







    def       UpperCamelCase__       (     self)	->     List[str]:
        __UpperCamelCase					:List[str]        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Any        =		create_components()
        __UpperCamelCase					:Dict        =		[1, 2, 3]

        # This should work
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Tuple        =		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       UpperCamelCase__       (     self)	->     int:
        from transformers import AutoModelForCausalLM

        __UpperCamelCase					:Optional[int]        =		AutoModelForCausalLM.from_pretrained(
            '''EleutherAI/gpt-neo-125m'''			,			load_in_abit=__lowercase			,			device_map={'''''': 0}			,			)
        __UpperCamelCase					:Optional[Any]        =		Accelerator()

        # This should work
        __UpperCamelCase					:int        =		accelerator.prepare(__lowercase)







    @slow
    @require_bnb
    def       UpperCamelCase__       (     self)	->     List[str]:
        from transformers import AutoModelForCausalLM

        __UpperCamelCase					:str        =		Accelerator()

        with init_empty_weights():
            __UpperCamelCase					:List[Any]        =		AutoModelForCausalLM.from_pretrained(
                '''EleutherAI/gpt-neo-125m'''			,			)
            model.tie_weights()
            __UpperCamelCase					:List[str]        =		infer_auto_device_map(__lowercase)
            __UpperCamelCase					:str        =		'''cpu'''

        __UpperCamelCase					:List[Any]        =		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):
            __UpperCamelCase					:Union[str, Any]        =		accelerator.prepare(__lowercase)







    @slow
    @require_bnb
    @require_multi_gpu
    def       UpperCamelCase__       (     self)	->     Union[str, Any]:
        from transformers import AutoModelForCausalLM

        __UpperCamelCase					:int        =		{'''distributed_type''': DistributedType.MULTI_GPU}

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

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

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

        PartialState._reset_state()







    @slow
    @require_bnb
    @require_multi_gpu
    def       UpperCamelCase__       (     self)	->     Dict:
        from transformers import AutoModelForCausalLM

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

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

        # This should work
        __UpperCamelCase					:int        =		accelerator.prepare(__lowercase)







    @require_cuda
    def       UpperCamelCase__       (     self)	->     int:
        __UpperCamelCase					:Tuple        =		torch.nn.Linear(10			,			10)
        __UpperCamelCase					:Optional[Any]        =		torch.optim.SGD(model.parameters()			,			lr=0.01)
        __UpperCamelCase					:Any        =		Accelerator(cpu=__lowercase)
        __UpperCamelCase					:Tuple        =		accelerator.prepare(__lowercase)





