

'''simple docstring'''






import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu






class      lowercase_      (nn.Module ):




  """simple docstring"""




  def __init__( self  :	Tuple     ):
   super().__init__()
   __lowercase				    =					nn.Linear(3   ,4     )
   __lowercase				    =					nn.BatchNormad(4     )
   __lowercase				    =					nn.Linear(4   ,5     )






  def    SCREAMING_SNAKE_CASE ( self  :	List[str]   ,lowercase__  :	List[Any]     ):
   return self.lineara(self.batchnorm(self.lineara(lowercase__     )     )     )






class      lowercase_      (lowerCamelCase__ ):




  """simple docstring"""




  def    SCREAMING_SNAKE_CASE ( self  :	List[Any]   ,lowercase__  :	Dict   ,*lowercase__  :	List[str]   ,**lowercase__  :	List[str]     ):
   return (args[0] + 1,) + args[1:], kwargs






class      lowercase_      (lowerCamelCase__ ):




  """simple docstring"""




  def    SCREAMING_SNAKE_CASE ( self  :	List[str]   ,lowercase__  :	Optional[Any]   ,lowercase__  :	Optional[Any]     ):
   return output + 1






class      lowercase_      (unittest.TestCase ):




  """simple docstring"""




  def    SCREAMING_SNAKE_CASE ( self  :	str     ):
   __lowercase				    =					ModelForTest()
   __lowercase				    =					ModelHook()

   add_hook_to_module(lowercase__   ,lowercase__     )
   self.assertEqual(test_model._hf_hook   ,lowercase__     )
   self.assertTrue(hasattr(lowercase__   ,'''_old_forward'''     )     )

   # Check adding the hook did not change the name or the signature
   self.assertEqual(test_model.forward.__name__   ,'''forward'''     )
   self.assertListEqual(list(inspect.signature(test_model.forward     ).parameters     )   ,['''x''']     )

   remove_hook_from_module(lowercase__     )
   self.assertFalse(hasattr(lowercase__   ,'''_hf_hook'''     )     )
   self.assertFalse(hasattr(lowercase__   ,'''_old_forward'''     )     )






  def    SCREAMING_SNAKE_CASE ( self  :	List[str]     ):
   __lowercase				    =					ModelForTest()
   __lowercase				    =					ModelHook()

   add_hook_to_module(lowercase__   ,lowercase__     )
   add_hook_to_module(lowercase__   ,lowercase__   ,append=lowercase__     )

   self.assertEqual(isinstance(test_model._hf_hook   ,lowercase__     )   ,lowercase__     )
   self.assertEqual(len(test_model._hf_hook.hooks     )   ,2     )
   self.assertTrue(hasattr(lowercase__   ,'''_old_forward'''     )     )

   # Check adding the hook did not change the name or the signature
   self.assertEqual(test_model.forward.__name__   ,'''forward'''     )
   self.assertListEqual(list(inspect.signature(test_model.forward     ).parameters     )   ,['''x''']     )

   remove_hook_from_module(lowercase__     )
   self.assertFalse(hasattr(lowercase__   ,'''_hf_hook'''     )     )
   self.assertFalse(hasattr(lowercase__   ,'''_old_forward'''     )     )






  def    SCREAMING_SNAKE_CASE ( self  :	Optional[int]     ):
   __lowercase				    =					ModelForTest()
   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					test_model(x + 1     )
   __lowercase				    =					test_model(x + 2     )

   __lowercase				    =					PreForwardHook()
   add_hook_to_module(lowercase__   ,lowercase__     )
   __lowercase				    =					test_model(lowercase__     )
   self.assertTrue(torch.allclose(lowercase__   ,lowercase__   ,atol=1e-5     )     )

   # Attaching a hook to a model when it already has one replaces, does not chain
   __lowercase				    =					PreForwardHook()
   add_hook_to_module(lowercase__   ,lowercase__     )
   __lowercase				    =					test_model(lowercase__     )
   self.assertTrue(torch.allclose(lowercase__   ,lowercase__   ,atol=1e-5     )     )

   # You need to use the sequential hook to chain two or more hooks
   __lowercase				    =					SequentialHook(PreForwardHook()   ,PreForwardHook()     )
   add_hook_to_module(lowercase__   ,lowercase__     )

   __lowercase				    =					test_model(lowercase__     )
   assert torch.allclose(lowercase__   ,lowercase__   ,atol=1e-5     )






  def    SCREAMING_SNAKE_CASE ( self  :	Optional[int]     ):
   __lowercase				    =					ModelForTest()
   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					test_model(lowercase__     )

   __lowercase				    =					PostForwardHook()
   add_hook_to_module(lowercase__   ,lowercase__     )
   __lowercase				    =					test_model(lowercase__     )
   self.assertTrue(torch.allclose(lowercase__   ,output + 1   ,atol=1e-5     )     )

   # Attaching a hook to a model when it already has one replaces, does not chain
   __lowercase				    =					PostForwardHook()
   add_hook_to_module(lowercase__   ,lowercase__     )
   __lowercase				    =					test_model(lowercase__     )
   self.assertTrue(torch.allclose(lowercase__   ,output + 1   ,atol=1e-5     )     )

   # You need to use the sequential hook to chain two or more hooks
   __lowercase				    =					SequentialHook(PostForwardHook()   ,PostForwardHook()     )
   add_hook_to_module(lowercase__   ,lowercase__     )

   __lowercase				    =					test_model(lowercase__     )
   assert torch.allclose(lowercase__   ,output + 2   ,atol=1e-5     )






  def    SCREAMING_SNAKE_CASE ( self  :	int     ):
   __lowercase				    =					ModelForTest()
   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					test_model(lowercase__     )

   __lowercase				    =					PostForwardHook()
   add_hook_to_module(lowercase__   ,lowercase__     )
   __lowercase				    =					test_model(lowercase__     )
   self.assertTrue(torch.allclose(lowercase__   ,output + 1     )     )
   self.assertTrue(outputa.requires_grad     )

   __lowercase				    =					True
   __lowercase				    =					test_model(lowercase__     )
   self.assertFalse(outputa.requires_grad     )






  @require_multi_gpu
  def    SCREAMING_SNAKE_CASE ( self  :	Dict     ):
   __lowercase				    =					ModelForTest()
   # Everything is on CPU
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )

   # This will move each submodule on different devices
   add_hook_to_module(model.lineara   ,AlignDevicesHook(execution_device=0     )     )
   add_hook_to_module(model.batchnorm   ,AlignDevicesHook(execution_device=0     )     )
   add_hook_to_module(model.lineara   ,AlignDevicesHook(execution_device=1     )     )

   self.assertEqual(model.lineara.weight.device   ,torch.device(0     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device(0     )     )
   self.assertEqual(model.batchnorm.running_mean.device   ,torch.device(0     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device(1     )     )

   # We can still make a forward pass. The input does not need to be on any particular device
   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					model(lowercase__     )
   self.assertEqual(output.device   ,torch.device(1     )     )

   # We can add a general hook to put back output on same device as input.
   add_hook_to_module(lowercase__   ,AlignDevicesHook(io_same_device=lowercase__     )     )
   __lowercase				    =					torch.randn(2   ,3     ).to(0     )
   __lowercase				    =					model(lowercase__     )
   self.assertEqual(output.device   ,torch.device(0     )     )






  def    SCREAMING_SNAKE_CASE ( self  :	str     ):
   __lowercase				    =					ModelForTest()

   # Everything is on CPU
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )

   # This will move each submodule on different devices
   __lowercase				    =					{'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

   add_hook_to_module(model.lineara   ,AlignDevicesHook(**lowercase__     )     )
   add_hook_to_module(model.batchnorm   ,AlignDevicesHook(**lowercase__     )     )
   add_hook_to_module(model.lineara   ,AlignDevicesHook(**lowercase__     )     )

   # Parameters have been offloaded, so on the meta device
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   # Buffers are not included in the offload by default, so are on the execution device
   __lowercase				    =					torch.device(hook_kwargs['''execution_device''']     )
   self.assertEqual(model.batchnorm.running_mean.device   ,lowercase__     )

   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					model(lowercase__     )
   self.assertEqual(output.device   ,lowercase__     )

   # Removing hooks loads back the weights in the model.
   remove_hook_from_module(model.lineara     )
   remove_hook_from_module(model.batchnorm     )
   remove_hook_from_module(model.lineara     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )

   # Now test with buffers included in the offload
   __lowercase				    =					{
       '''execution_device''': 0 if torch.cuda.is_available() else '''cpu''',
       '''offload''': True,
       '''offload_buffers''': True,
   }

   add_hook_to_module(model.lineara   ,AlignDevicesHook(**lowercase__     )     )
   add_hook_to_module(model.batchnorm   ,AlignDevicesHook(**lowercase__     )     )
   add_hook_to_module(model.lineara   ,AlignDevicesHook(**lowercase__     )     )

   # Parameters have been offloaded, so on the meta device, buffers included
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.running_mean.device   ,torch.device('''meta'''     )     )

   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					model(lowercase__     )
   self.assertEqual(output.device   ,lowercase__     )

   # Removing hooks loads back the weights in the model.
   remove_hook_from_module(model.lineara     )
   remove_hook_from_module(model.batchnorm     )
   remove_hook_from_module(model.lineara     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )






  def    SCREAMING_SNAKE_CASE ( self  :	int     ):
   __lowercase				    =					ModelForTest()

   # Everything is on CPU
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )

   # This will move each submodule on different devices
   __lowercase				    =					0 if torch.cuda.is_available() else '''cpu'''
   attach_align_device_hook(lowercase__   ,execution_device=lowercase__   ,offload=lowercase__     )

   # Parameters have been offloaded, so on the meta device
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   # Buffers are not included in the offload by default, so are on the execution device
   __lowercase				    =					torch.device(lowercase__     )
   self.assertEqual(model.batchnorm.running_mean.device   ,lowercase__     )

   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					model(lowercase__     )
   self.assertEqual(output.device   ,lowercase__     )

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(lowercase__     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )

   # Now test with buffers included in the offload
   attach_align_device_hook(lowercase__   ,execution_device=lowercase__   ,offload=lowercase__   ,offload_buffers=lowercase__     )

   # Parameters have been offloaded, so on the meta device, buffers included
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.running_mean.device   ,torch.device('''meta'''     )     )

   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					model(lowercase__     )
   self.assertEqual(output.device   ,lowercase__     )

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(lowercase__     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )






  def    SCREAMING_SNAKE_CASE ( self  :	List[Any]     ):
   __lowercase				    =					ModelForTest()

   # Everything is on CPU
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )

   # This will move each submodule on different devices
   __lowercase				    =					0 if torch.cuda.is_available() else '''cpu'''
   attach_align_device_hook(
       lowercase__   ,execution_device=lowercase__   ,offload=lowercase__   ,weights_map=model.state_dict()     )

   # Parameters have been offloaded, so on the meta device
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   # Buffers are not included in the offload by default, so are on the execution device
   __lowercase				    =					torch.device(lowercase__     )
   self.assertEqual(model.batchnorm.running_mean.device   ,lowercase__     )

   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					model(lowercase__     )
   self.assertEqual(output.device   ,lowercase__     )

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(lowercase__     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )

   # Now test with buffers included in the offload
   attach_align_device_hook(
       lowercase__   ,execution_device=lowercase__   ,offload=lowercase__   ,weights_map=model.state_dict()   ,offload_buffers=lowercase__   ,)

   # Parameters have been offloaded, so on the meta device, buffers included
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''meta'''     )     )
   self.assertEqual(model.batchnorm.running_mean.device   ,torch.device('''meta'''     )     )

   __lowercase				    =					torch.randn(2   ,3     )
   __lowercase				    =					model(lowercase__     )
   self.assertEqual(output.device   ,lowercase__     )

   # Removing hooks loads back the weights in the model.
   remove_hook_from_submodules(lowercase__     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.batchnorm.weight.device   ,torch.device('''cpu'''     )     )
   self.assertEqual(model.lineara.weight.device   ,torch.device('''cpu'''     )     )




from math import isqrt, loga







def 						lowerCamelCase__      ( __A    :int		):


       """simple docstring"""




       __snake_case       	=							[True] * max_number
       for i in range(2	,isqrt(max_number - 1		) + 1		):
              if is_prime[i]:
                     for j in range(i**2	,__A	,__A		):
                            __snake_case       	=							False

       return [i for i in range(2	,__A		) if is_prime[i]]







def 						lowerCamelCase__      ( __A    :int = 8_0_0_8_0_0	,__A    :int = 8_0_0_8_0_0		):


       """simple docstring"""




       __snake_case       	=							degree * loga(__A		)
       __snake_case       	=							int(__A		)
       __snake_case       	=							calculate_prime_numbers(__A		)

       __snake_case       	=							0
       __snake_case       	=							0
       __snake_case       	=							len(__A		) - 1
       while left < right:
              while (
                  prime_numbers[right] * loga(prime_numbers[left]		)
                  + prime_numbers[left] * loga(prime_numbers[right]		)
                  > upper_bound
              ):
                     right -= 1
              hybrid_integers_count += right - left
              left += 1

       return hybrid_integers_count


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

