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

  return model, optimizer, scheduler, train_dl, valid_dl


def 						_UpperCAmelCase					(	_lowerCamelCase    :							Optional[int]		)     ->    int:
  return (model.weight.abs().sum() + model.bias.abs().sum()).item()


def 						_UpperCAmelCase					(	_lowerCamelCase    :							Tuple		)     ->    List[str]:
  _lowerCAmelCase							:	List[str] 					=       torch.nn.Linear(*tuple(model.weight.T.shape		)		).state_dict()
  model.load_state_dict(lowercase__		)



class 	a_	(_a     ):




       @require_cuda
       def 				__UpperCamelCase							(  self	):
         _lowerCAmelCase							:	Tuple 					=       Accelerator()
         assert PartialState._shared_state["_cpu"] is False
         assert PartialState._shared_state["device"].type == "cuda"
         with self.assertRaises(__lowercase	):
           _lowerCAmelCase							:	List[str] 					=       Accelerator(cpu=__lowercase	)




       def 				__UpperCamelCase							(  self	):
         _lowerCAmelCase							:	Union[str, Any] 					=       Accelerator()
         _lowerCAmelCase							:	List[Any] 					=       GradientState()
         assert state.num_steps == 1
         _lowerCAmelCase							:	Optional[Any] 					=       4
         assert state.num_steps == 4

         assert state.sync_gradients is True
         _lowerCAmelCase							:	Optional[int] 					=       False
         assert state.sync_gradients is False
         GradientState._reset_state()




       def 				__UpperCamelCase							(  self	):
         _lowerCAmelCase							:	Optional[Any] 					=       Accelerator()
         _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase							:	int 					=       create_components()

         (
             (
             _lowerCAmelCase  
         )					,       (
             _lowerCAmelCase  
         )					,       (
             _lowerCAmelCase  
         )					,       (
             _lowerCAmelCase  
         )					,       (
             _lowerCAmelCase  
         )					,         
         )							:	Any 					=       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	):
         _lowerCAmelCase							:	Union[str, Any] 					=       Accelerator()
         _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase							:	Union[str, 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	):
         PartialState._reset_state()

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

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




       def 				__UpperCamelCase							(  self	):
         _lowerCAmelCase							:	List[str] 					=       Accelerator()
         _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase							:	Tuple 					=       create_components()
         accelerator.prepare(__lowercase , __lowercase , __lowercase , __lowercase , __lowercase	)

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

         _lowerCAmelCase							:	Dict 					=       get_signature(__lowercase	)

         # saving hook
         def save_config(snake_case_ , snake_case_ , snake_case_	):
           _lowerCAmelCase							:	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(snake_case_ , snake_case_	):
           with open(os.path.join(__lowercase , """data.json"""	) , """r"""	) as f:
             _lowerCAmelCase							:	Any 					=       json.load(__lowercase	)

           _lowerCAmelCase							:	List[Any] 					=       config["""class_name"""]

         _lowerCAmelCase							:	Tuple 					=       accelerator.register_save_state_pre_hook(__lowercase	)
         _lowerCAmelCase							:	str 					=       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
           _lowerCAmelCase							:	Optional[Any] 					=       """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
           _lowerCAmelCase							:	Optional[Any] 					=       """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	):
         _lowerCAmelCase							:	List[str] 					=       Accelerator()
         _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase							:	Union[str, Any] 					=       create_components()
         _lowerCAmelCase							:	str 					=       None

         # This should work
         _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase							:	Optional[int] 					=       accelerator.prepare(
             __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase	)
         self.assertTrue(dummy_obj is None	)




       def 				__UpperCamelCase							(  self	):
         _lowerCAmelCase							:	Optional[int] 					=       Accelerator()
         _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase							:	int 					=       create_components()
         _lowerCAmelCase							:	str 					=       [1, 2, 3]

         # This should work
         _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase					,       _lowerCAmelCase							:	Optional[Any] 					=       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	):
         from transformers import AutoModelForCausalLM

         _lowerCAmelCase							:	Tuple 					=       AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m""" , load_in_abit=__lowercase , device_map={"""""": 0} , )
         _lowerCAmelCase							:	Union[str, Any] 					=       Accelerator()

         # This should work
         _lowerCAmelCase							:	str 					=       accelerator.prepare(__lowercase	)




       @slow
       @require_bnb
       def 				__UpperCamelCase							(  self	):
         from transformers import AutoModelForCausalLM

         _lowerCAmelCase							:	List[Any] 					=       Accelerator()

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

         _lowerCAmelCase							:	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	):
           _lowerCAmelCase							:	str 					=       accelerator.prepare(__lowercase	)




       @slow
       @require_bnb
       @require_multi_gpu
       def 				__UpperCamelCase							(  self	):
         from transformers import AutoModelForCausalLM

         _lowerCAmelCase							:	str 					=       {"""distributed_type""": DistributedType.MULTI_GPU}

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

         _lowerCAmelCase							:	List[str] 					=       AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m""" , load_in_abit=__lowercase , device_map=__lowercase , )
         _lowerCAmelCase							:	Optional[Any] 					=       Accelerator()

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

         PartialState._reset_state()




       @slow
       @require_bnb
       @require_multi_gpu
       def 				__UpperCamelCase							(  self	):
         from transformers import AutoModelForCausalLM

         with init_empty_weights():
           _lowerCAmelCase							:	Any 					=       AutoModelForCausalLM.from_pretrained(
               """EleutherAI/gpt-neo-125m""" , )
           _lowerCAmelCase							:	Dict 					=       infer_auto_device_map(__lowercase	)
           _lowerCAmelCase							:	Union[str, Any] 					=       1

         _lowerCAmelCase							:	Optional[Any] 					=       AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m""" , load_in_abit=__lowercase , device_map=__lowercase , )
         _lowerCAmelCase							:	Optional[Any] 					=       Accelerator()

         # This should work
         _lowerCAmelCase							:	Optional[Any] 					=       accelerator.prepare(__lowercase	)







       @require_cuda
       def 				__UpperCamelCase							(  self	):
         _lowerCAmelCase							:	Union[str, Any] 					=       torch.nn.Linear(1_0 , 1_0	)
         _lowerCAmelCase							:	int 					=       torch.optim.SGD(model.parameters() , lr=0.01	)
         _lowerCAmelCase							:	Any 					=       Accelerator(cpu=__lowercase	)
         _lowerCAmelCase							:	Optional[int] 					=       accelerator.prepare(__lowercase	)





'''simple docstring'''

from collections import defaultdict
from math import gcd

def 					__UpperCamelCase				(						lowercase__    :	int = 1_50_00_00			):






     '''simple docstring'''





     __lowercase      		=defaultdict(lowercase__			)
     __lowercase      		=2
     while 2 * euclid_m * (euclid_m + 1) <= limit:
          for euclid_n in range((euclid_m % 2) + 1,       lowercase__,       2			):
               if gcd(lowercase__,       lowercase__			) > 1:
                    continue
               __lowercase      		=2 * euclid_m * (euclid_m + euclid_n)
               for perimeter in range(lowercase__,       limit + 1,       lowercase__			):
                    frequencies[perimeter] += 1
          euclid_m += 1
     return sum(1 for frequency in frequencies.values() if frequency == 1			)


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


