


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        lowerCamelCase      (     nn.Module   ):


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


	def    A(						self							,    lowercase__):
							return self.lineara(self.batchnorm(self.lineara(lowercase__)))






class        lowerCamelCase      (     _UpperCamelCase   ):


	def    A(						self							,    lowercase__							,    *lowercase__							,    **lowercase__):
							return (args[0] + 1,) + args[1:], kwargs






class        lowerCamelCase      (     _UpperCamelCase   ):


	def    A(						self							,    lowercase__							,    lowercase__):
							return output + 1






class        lowerCamelCase      (     unittest.TestCase   ):


	def    A(						self):
							__UpperCAmelCase :       Union[str, Any]  						=  ModelForTest()
							__UpperCAmelCase :       str  						=  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    A(						self):
							__UpperCAmelCase :       Optional[int]  						=  ModelForTest()
							__UpperCAmelCase :       str  						=  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    A(						self):
							__UpperCAmelCase :       Tuple  						=  ModelForTest()
							__UpperCAmelCase :       Union[str, Any]  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       str  						=  test_model(x + 1)
							__UpperCAmelCase :       List[Any]  						=  test_model(x + 2)

							__UpperCAmelCase :       Union[str, Any]  						=  PreForwardHook()
							add_hook_to_module(lowercase__							,    lowercase__)
							__UpperCAmelCase :       List[Any]  						=  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
							__UpperCAmelCase :       Union[str, Any]  						=  PreForwardHook()
							add_hook_to_module(lowercase__							,    lowercase__)
							__UpperCAmelCase :       List[Any]  						=  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
							__UpperCAmelCase :       Optional[Any]  						=  SequentialHook(PreForwardHook()							,    PreForwardHook())
							add_hook_to_module(lowercase__							,    lowercase__)

							__UpperCAmelCase :       Optional[int]  						=  test_model(lowercase__)
							assert torch.allclose(lowercase__							,    lowercase__							,    atol=1e-5)


	def    A(						self):
							__UpperCAmelCase :       Optional[int]  						=  ModelForTest()
							__UpperCAmelCase :       Dict  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       List[str]  						=  test_model(lowercase__)

							__UpperCAmelCase :       Optional[Any]  						=  PostForwardHook()
							add_hook_to_module(lowercase__							,    lowercase__)
							__UpperCAmelCase :       str  						=  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
							__UpperCAmelCase :       Optional[int]  						=  PostForwardHook()
							add_hook_to_module(lowercase__							,    lowercase__)
							__UpperCAmelCase :       Any  						=  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
							__UpperCAmelCase :       Any  						=  SequentialHook(PostForwardHook()							,    PostForwardHook())
							add_hook_to_module(lowercase__							,    lowercase__)

							__UpperCAmelCase :       List[Any]  						=  test_model(lowercase__)
							assert torch.allclose(lowercase__							,    output + 2							,    atol=1e-5)


	def    A(						self):
							__UpperCAmelCase :       Optional[int]  						=  ModelForTest()
							__UpperCAmelCase :       Any  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       Optional[Any]  						=  test_model(lowercase__)

							__UpperCAmelCase :       Tuple  						=  PostForwardHook()
							add_hook_to_module(lowercase__							,    lowercase__)
							__UpperCAmelCase :       Optional[Any]  						=  test_model(lowercase__)
							self.assertTrue(torch.allclose(lowercase__							,    output + 1))
							self.assertTrue(outputa.requires_grad)

							__UpperCAmelCase :       Optional[int]  						=  True
							__UpperCAmelCase :       Dict  						=  test_model(lowercase__)
							self.assertFalse(outputa.requires_grad)


	@require_multi_gpu
	def    A(						self):
							__UpperCAmelCase :       str  						=  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
							__UpperCAmelCase :       Any  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       Union[str, Any]  						=  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__))
							__UpperCAmelCase :       int  						=  torch.randn(2							,    3).to(0)
							__UpperCAmelCase :       str  						=  model(lowercase__)
							self.assertEqual(output.device							,    torch.device(0))


	def    A(						self):
							__UpperCAmelCase :       List[Any]  						=  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
							__UpperCAmelCase :       List[str]  						=  {'''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
							__UpperCAmelCase :       Dict  						=  torch.device(hook_kwargs['''execution_device'''])
							self.assertEqual(model.batchnorm.running_mean.device							,    lowercase__)

							__UpperCAmelCase :       str  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       int  						=  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
							__UpperCAmelCase :       str  						=  {
							    '''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'''))

							__UpperCAmelCase :       Dict  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       Optional[int]  						=  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    A(						self):
							__UpperCAmelCase :       Any  						=  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
							__UpperCAmelCase :       str  						=  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
							__UpperCAmelCase :       Union[str, Any]  						=  torch.device(lowercase__)
							self.assertEqual(model.batchnorm.running_mean.device							,    lowercase__)

							__UpperCAmelCase :       Optional[Any]  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       Any  						=  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'''))

							__UpperCAmelCase :       List[str]  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       int  						=  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    A(						self):
							__UpperCAmelCase :       Dict  						=  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
							__UpperCAmelCase :       Tuple  						=  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
							__UpperCAmelCase :       Optional[int]  						=  torch.device(lowercase__)
							self.assertEqual(model.batchnorm.running_mean.device							,    lowercase__)

							__UpperCAmelCase :       Optional[int]  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       str  						=  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'''))

							__UpperCAmelCase :       Dict  						=  torch.randn(2							,    3)
							__UpperCAmelCase :       Dict  						=  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       (					lowercase_			,					lowercase_			)		->    str:




						'''simple docstring'''







						if a < 0 or b < 0:
												raise ValueError('''the value of both inputs must be positive'''			)
						__UpperCAmelCase :       Dict  						=  str(bin(lowercase_			)			)[2:]  # remove the leading "0b"
						__UpperCAmelCase :       List[Any]  						=  str(bin(lowercase_			)			)[2:]
						__UpperCAmelCase :       List[Any]  						=  max(len(lowercase_			)			,					len(lowercase_			)			)
						return "0b" + "".join(
						    str(int('''1''' in (char_a, char_b)			)			)
						    for char_a, char_b in zip(a_binary.zfill(lowercase_			)			,					b_binary.zfill(lowercase_			)			)			)


if __name__ == "__main__":
				import doctest

				doctest.testmod()

