


'''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       (		)						->					List[Any]:



       '''simple docstring'''





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

       return model, optimizer, scheduler, train_dl, valid_dl





def __UpperCAmelCase       (		__A			)						->					Union[str, Any]:



       '''simple docstring'''





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





def __UpperCAmelCase       (		__A			)						->					List[str]:



       '''simple docstring'''





       UpperCAmelCase__										=     torch.nn.Linear(*tuple(model.weight.T.shape			)			).state_dict()
       model.load_state_dict(__A			)




class     lowercase__  (					__SCREAMING_SNAKE_CASE		):






  @require_cuda
  def 			_UpperCAmelCase  (					self      :					List[Any] ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					Optional[int] ):




         """simple docstring"""







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

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






  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         UpperCAmelCase__										=     Accelerator()
         UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__										=     create_components()

         (
             (
             UpperCAmelCase__			
         )						,			(
             UpperCAmelCase__			
         )						,			(
             UpperCAmelCase__			
         )						,			(
             UpperCAmelCase__			
         )						,			(
             UpperCAmelCase__			
         )						,						
         )										=     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      :					Dict ):




         """simple docstring"""







         UpperCAmelCase__										=     Accelerator()
         UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__										=     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      :					int ):




         """simple docstring"""







         PartialState._reset_state()

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

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






  def 			_UpperCAmelCase  (					self      :					str ):




         """simple docstring"""







         UpperCAmelCase__										=     Accelerator()
         UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__										=     create_components()
         accelerator.prepare(_lowercase		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase )

         UpperCAmelCase__										=     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 ):




         """simple docstring"""







         UpperCAmelCase__										=     Accelerator()
         UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__										=     create_components()
         accelerator.prepare(_lowercase		,						_lowercase		,						_lowercase		,						_lowercase		,						_lowercase )

         UpperCAmelCase__										=     get_signature(_lowercase )

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

                UpperCAmelCase__										=     config["class_name"]

         UpperCAmelCase__										=     accelerator.register_save_state_pre_hook(_lowercase )
         UpperCAmelCase__										=     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__										=     "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__										=     "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      :					Union[str, Any] ):




         """simple docstring"""







         UpperCAmelCase__										=     Accelerator()
         UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__										=     create_components()
         UpperCAmelCase__										=     None

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






  def 			_UpperCAmelCase  (					self      :					Tuple ):




         """simple docstring"""







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

         # This should work
         UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__						,			UpperCAmelCase__										=     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      :					Optional[int] ):




         """simple docstring"""







         from transformers import AutoModelForCausalLM

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

         # This should work
         UpperCAmelCase__										=     accelerator.prepare(_lowercase )






  @slow
  @require_bnb
  def 			_UpperCAmelCase  (					self      :					List[Any] ):




         """simple docstring"""







         from transformers import AutoModelForCausalLM

         UpperCAmelCase__										=     Accelerator()

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

         UpperCAmelCase__										=     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__										=     accelerator.prepare(_lowercase )






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




         """simple docstring"""







         from transformers import AutoModelForCausalLM

         UpperCAmelCase__										=     {"distributed_type": DistributedType.MULTI_GPU}

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

         UpperCAmelCase__										=     AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"		,						load_in_abit=_lowercase		,						device_map=_lowercase		,						)
         UpperCAmelCase__										=     Accelerator()

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

         PartialState._reset_state()






  @slow
  @require_bnb
  @require_multi_gpu
  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         from transformers import AutoModelForCausalLM

         with init_empty_weights():
                UpperCAmelCase__										=     AutoModelForCausalLM.from_pretrained(
                    "EleutherAI/gpt-neo-125m"		,						)
                UpperCAmelCase__										=     infer_auto_device_map(_lowercase )
                UpperCAmelCase__										=     1

         UpperCAmelCase__										=     AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"		,						load_in_abit=_lowercase		,						device_map=_lowercase		,						)
         UpperCAmelCase__										=     Accelerator()

         # This should work
         UpperCAmelCase__										=     accelerator.prepare(_lowercase )






  @require_cuda
  def 			_UpperCAmelCase  (					self      :					Optional[int] ):




         """simple docstring"""







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









from __future__ import annotations





def __UpperCAmelCase       (		__A			)						->					list[int]:



       '''simple docstring'''





       UpperCAmelCase__										=     [True] * limit
       UpperCAmelCase__										=     False
       UpperCAmelCase__										=     False
       UpperCAmelCase__										=     True

       for i in range(3 ,						int(limit**0.5 + 1			) ,						2			):
              UpperCAmelCase__										=     i * 2
              while index < limit:
                     UpperCAmelCase__										=     False
                     UpperCAmelCase__										=     index + i

       UpperCAmelCase__										=     [2]

       for i in range(3 ,						__A ,						2			):
              if is_prime[i]:
                     primes.append(__A			)

       return primes





def __UpperCAmelCase       (		__A = 1_0_0_0_0_0_0			)						->					int:



       '''simple docstring'''





       UpperCAmelCase__										=     prime_sieve(__A			)
       UpperCAmelCase__										=     0
       UpperCAmelCase__										=     0

       for i in range(len(__A			)			):
              for j in range(i + length ,						len(__A			)			):
                     UpperCAmelCase__										=     sum(primes[i:j]			)
                     if sol >= ceiling:
                            break

                     if sol in primes:
                            UpperCAmelCase__										=     j - i
                            UpperCAmelCase__										=     sol

       return largest


if __name__ == "__main__":
    print(f"{solution() = }")





