




"""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     ):
								super().__init__()
								__lowerCamelCase					:	Tuple	       =				nn.Linear(3						,	4     )
								__lowerCamelCase					:	Optional[Any]	       =				nn.BatchNormad(4     )
								__lowerCamelCase					:	Optional[Any]	       =				nn.Linear(4						,	5     )





					def     snake_case_				(    self						,	__a     ):
								return self.lineara(self.batchnorm(self.lineara(__a     )     )     )



class __lowercase(      __lowercase				):




					'''simple docstring'''


					def     snake_case_				(    self						,	__a						,	*__a						,	**__a     ):
								return (args[0] + 1,) + args[1:], kwargs



class __lowercase(      __lowercase				):




					'''simple docstring'''


					def     snake_case_				(    self						,	__a						,	__a     ):
								return output + 1



class __lowercase(      unittest.TestCase				):




					'''simple docstring'''


					def     snake_case_				(    self     ):
								__lowerCamelCase					:	List[str]	       =				ModelForTest()
								__lowerCamelCase					:	Dict	       =				ModelHook()

								add_hook_to_module(__a						,	__a     )
								self.assertEqual(test_model._hf_hook						,	__a     )
								self.assertTrue(hasattr(__a						,	'_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(__a     )
								self.assertFalse(hasattr(__a						,	'_hf_hook'     )     )
								self.assertFalse(hasattr(__a						,	'_old_forward'     )     )
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	List[str]	       =				ModelForTest()
								__lowerCamelCase					:	Union[str, Any]	       =				ModelHook()

								add_hook_to_module(__a						,	__a     )
								add_hook_to_module(__a						,	__a						,	append=__a     )

								self.assertEqual(isinstance(test_model._hf_hook						,	__a     )						,	__a     )
								self.assertEqual(len(test_model._hf_hook.hooks     )						,	2     )
								self.assertTrue(hasattr(__a						,	'_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(__a     )
								self.assertFalse(hasattr(__a						,	'_hf_hook'     )     )
								self.assertFalse(hasattr(__a						,	'_old_forward'     )     )
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	str	       =				ModelForTest()
								__lowerCamelCase					:	Dict	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	Union[str, Any]	       =				test_model(x + 1     )
								__lowerCamelCase					:	Optional[int]	       =				test_model(x + 2     )

								__lowerCamelCase					:	List[Any]	       =				PreForwardHook()
								add_hook_to_module(__a						,	__a     )
								__lowerCamelCase					:	Optional[int]	       =				test_model(__a     )
								self.assertTrue(torch.allclose(__a						,	__a						,	atol=1E-5     )     )

								# Attaching a hook to a model when it already has one replaces, does not chain
								__lowerCamelCase					:	Dict	       =				PreForwardHook()
								add_hook_to_module(__a						,	__a     )
								__lowerCamelCase					:	Tuple	       =				test_model(__a     )
								self.assertTrue(torch.allclose(__a						,	__a						,	atol=1E-5     )     )

								# You need to use the sequential hook to chain two or more hooks
								__lowerCamelCase					:	Any	       =				SequentialHook(PreForwardHook()						,	PreForwardHook()     )
								add_hook_to_module(__a						,	__a     )

								__lowerCamelCase					:	Optional[Any]	       =				test_model(__a     )
								assert torch.allclose(__a						,	__a						,	atol=1E-5     )
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	str	       =				ModelForTest()
								__lowerCamelCase					:	List[str]	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	int	       =				test_model(__a     )

								__lowerCamelCase					:	Dict	       =				PostForwardHook()
								add_hook_to_module(__a						,	__a     )
								__lowerCamelCase					:	Tuple	       =				test_model(__a     )
								self.assertTrue(torch.allclose(__a						,	output + 1						,	atol=1E-5     )     )

								# Attaching a hook to a model when it already has one replaces, does not chain
								__lowerCamelCase					:	str	       =				PostForwardHook()
								add_hook_to_module(__a						,	__a     )
								__lowerCamelCase					:	Optional[Any]	       =				test_model(__a     )
								self.assertTrue(torch.allclose(__a						,	output + 1						,	atol=1E-5     )     )

								# You need to use the sequential hook to chain two or more hooks
								__lowerCamelCase					:	Union[str, Any]	       =				SequentialHook(PostForwardHook()						,	PostForwardHook()     )
								add_hook_to_module(__a						,	__a     )

								__lowerCamelCase					:	str	       =				test_model(__a     )
								assert torch.allclose(__a						,	output + 2						,	atol=1E-5     )
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	List[Any]	       =				ModelForTest()
								__lowerCamelCase					:	Optional[int]	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	List[str]	       =				test_model(__a     )

								__lowerCamelCase					:	Any	       =				PostForwardHook()
								add_hook_to_module(__a						,	__a     )
								__lowerCamelCase					:	str	       =				test_model(__a     )
								self.assertTrue(torch.allclose(__a						,	output + 1     )     )
								self.assertTrue(outputa.requires_grad     )

								__lowerCamelCase					:	Optional[int]	       =				True
								__lowerCamelCase					:	Optional[int]	       =				test_model(__a     )
								self.assertFalse(outputa.requires_grad     )
					@require_multi_gpu
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	int	       =				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
								__lowerCamelCase					:	str	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	Dict	       =				model(__a     )
								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(__a						,	AlignDevicesHook(io_same_device=__a     )     )
								__lowerCamelCase					:	Optional[int]	       =				torch.randn(2						,	3     ).to(0     )
								__lowerCamelCase					:	str	       =				model(__a     )
								self.assertEqual(output.device						,	torch.device(0     )     )
					def     snake_case_				(    self     ):
								__lowerCamelCase					:	Union[str, 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
								__lowerCamelCase					:	Tuple	       =				{"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

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

								# 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
								__lowerCamelCase					:	List[Any]	       =				torch.device(hook_kwargs['execution_device']     )
								self.assertEqual(model.batchnorm.running_mean.device						,	__a     )

								__lowerCamelCase					:	Optional[Any]	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	Optional[Any]	       =				model(__a     )
								self.assertEqual(output.device						,	__a     )

								# 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
								__lowerCamelCase					:	Any	       =				{
								    """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
								    """offload""": True,
								    """offload_buffers""": True,
								}

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

								# 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'     )     )

								__lowerCamelCase					:	int	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	Optional[int]	       =				model(__a     )
								self.assertEqual(output.device						,	__a     )

								# 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     snake_case_				(    self     ):
								__lowerCamelCase					:	int	       =				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
								__lowerCamelCase					:	int	       =				0 if torch.cuda.is_available() else """cpu"""
								attach_align_device_hook(__a						,	execution_device=__a						,	offload=__a     )

								# 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
								__lowerCamelCase					:	List[Any]	       =				torch.device(__a     )
								self.assertEqual(model.batchnorm.running_mean.device						,	__a     )

								__lowerCamelCase					:	Dict	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	Optional[Any]	       =				model(__a     )
								self.assertEqual(output.device						,	__a     )

								# Removing hooks loads back the weights in the model.
								remove_hook_from_submodules(__a     )
								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(__a						,	execution_device=__a						,	offload=__a						,	offload_buffers=__a     )

								# 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'     )     )

								__lowerCamelCase					:	Optional[int]	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	int	       =				model(__a     )
								self.assertEqual(output.device						,	__a     )

								# Removing hooks loads back the weights in the model.
								remove_hook_from_submodules(__a     )
								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     snake_case_				(    self     ):
								__lowerCamelCase					:	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
								__lowerCamelCase					:	Any	       =				0 if torch.cuda.is_available() else """cpu"""
								attach_align_device_hook(
								    __a						,	execution_device=__a						,	offload=__a						,	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
								__lowerCamelCase					:	List[Any]	       =				torch.device(__a     )
								self.assertEqual(model.batchnorm.running_mean.device						,	__a     )

								__lowerCamelCase					:	Dict	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	int	       =				model(__a     )
								self.assertEqual(output.device						,	__a     )

								# Removing hooks loads back the weights in the model.
								remove_hook_from_submodules(__a     )
								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(
								    __a						,	execution_device=__a						,	offload=__a						,	weights_map=model.state_dict()						,	offload_buffers=__a						,	)

								# 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'     )     )

								__lowerCamelCase					:	Tuple	       =				torch.randn(2						,	3     )
								__lowerCamelCase					:	Any	       =				model(__a     )
								self.assertEqual(output.device						,	__a     )

								# Removing hooks loads back the weights in the model.
								remove_hook_from_submodules(__a     )
								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'     )     )









"""simple docstring"""
a_				:    Dict										=     {
    '''meter''': '''m''',
    '''kilometer''': '''km''',
    '''megametre''': '''Mm''',
    '''gigametre''': '''Gm''',
    '''terametre''': '''Tm''',
    '''petametre''': '''Pm''',
    '''exametre''': '''Em''',
    '''zettametre''': '''Zm''',
    '''yottametre''': '''Ym''',
}
# Exponent of the factor(meter)
a_				:    Optional[int]										=     {
    '''m''': 0,
    '''km''': 3,
    '''Mm''': 6,
    '''Gm''': 9,
    '''Tm''': 12,
    '''Pm''': 15,
    '''Em''': 18,
    '''Zm''': 21,
    '''Ym''': 24,
}





def 						UpperCAmelCase				(			A__:    float							,				A__:    str							,				A__:    str     )  ->     float:
			__lowerCamelCase					:	List[Any]	       =				from_type.lower().strip('s'     )
			__lowerCamelCase					:	Optional[int]	       =				to_type.lower().strip('s'     )

			__lowerCamelCase					:	int	       =				UNIT_SYMBOL.get(A__							,				A__     )
			__lowerCamelCase					:	Dict	       =				UNIT_SYMBOL.get(A__							,				A__     )

			if from_sanitized not in METRIC_CONVERSION:
						__lowerCamelCase					:	Dict	       =				(
						    f'''Invalid \'from_type\' value: {from_type!r}.\n'''
						    f'''Conversion abbreviations are: {", ".join(A__     )}'''
						)
						raise ValueError(A__     )
			if to_sanitized not in METRIC_CONVERSION:
						__lowerCamelCase					:	Any	       =				(
						    f'''Invalid \'to_type\' value: {to_type!r}.\n'''
						    f'''Conversion abbreviations are: {", ".join(A__     )}'''
						)
						raise ValueError(A__     )
			__lowerCamelCase					:	Any	       =				METRIC_CONVERSION[from_sanitized]
			__lowerCamelCase					:	Tuple	       =				METRIC_CONVERSION[to_sanitized]
			__lowerCamelCase					:	Any	       =				1

			if from_exponent > to_exponent:
						__lowerCamelCase					:	Union[str, Any]	       =				from_exponent - to_exponent
			else:
						__lowerCamelCase					:	List[Any]	       =				-(to_exponent - from_exponent)

			return value * pow(10							,				A__     )


if __name__ == "__main__":
					from doctest import testmod

					testmod()



