

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







				def __init__(       self     :		Optional[int]   ):
										'''simple docstring'''
										super().__init__()
										lowercase :							Optional[int]    		=nn.Linear(3		,					4   )
										lowercase :							Optional[Any]    		=nn.BatchNormad(4   )
										lowercase :							Optional[int]    		=nn.Linear(4		,					5   )







				def  lowerCamelCase_	(       self     :		str		,					UpperCAmelCase__     :		Any   ):
										'''simple docstring'''
										return self.lineara(self.batchnorm(self.lineara(__A   )   )   )





class 		__SCREAMING_SNAKE_CASE							(			_lowercase      ):







				def  lowerCamelCase_	(       self     :		Union[str, Any]		,					UpperCAmelCase__     :		int		,					*UpperCAmelCase__     :		List[Any]		,					**UpperCAmelCase__     :		Any   ):
										'''simple docstring'''
										return (args[0] + 1,) + args[1:], kwargs





class 		__SCREAMING_SNAKE_CASE							(			_lowercase      ):







				def  lowerCamelCase_	(       self     :		Optional[int]		,					UpperCAmelCase__     :		List[str]		,					UpperCAmelCase__     :		str   ):
										'''simple docstring'''
										return output + 1





class 		__SCREAMING_SNAKE_CASE							(			unittest.TestCase      ):







				def  lowerCamelCase_	(       self     :		List[str]   ):
										'''simple docstring'''
										lowercase :							Union[str, Any]    		=ModelForTest()
										lowercase :							Union[str, Any]    		=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  lowerCamelCase_	(       self     :		Dict   ):
										'''simple docstring'''
										lowercase :							str    		=ModelForTest()
										lowercase :							Tuple    		=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  lowerCamelCase_	(       self     :		int   ):
										'''simple docstring'''
										lowercase :							int    		=ModelForTest()
										lowercase :							List[str]    		=torch.randn(2		,					3   )
										lowercase :							Tuple    		=test_model(x + 1   )
										lowercase :							List[str]    		=test_model(x + 2   )

										lowercase :							int    		=PreForwardHook()
										add_hook_to_module(__A		,					__A   )
										lowercase :							Union[str, Any]    		=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
										lowercase :							Dict    		=PreForwardHook()
										add_hook_to_module(__A		,					__A   )
										lowercase :							Union[str, Any]    		=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
										lowercase :							List[Any]    		=SequentialHook(PreForwardHook()		,					PreForwardHook()   )
										add_hook_to_module(__A		,					__A   )

										lowercase :							List[str]    		=test_model(__A   )
										assert torch.allclose(__A		,					__A		,					atol=1E-5   )







				def  lowerCamelCase_	(       self     :		Dict   ):
										'''simple docstring'''
										lowercase :							Any    		=ModelForTest()
										lowercase :							List[str]    		=torch.randn(2		,					3   )
										lowercase :							Union[str, Any]    		=test_model(__A   )

										lowercase :							Dict    		=PostForwardHook()
										add_hook_to_module(__A		,					__A   )
										lowercase :							int    		=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
										lowercase :							Optional[int]    		=PostForwardHook()
										add_hook_to_module(__A		,					__A   )
										lowercase :							int    		=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
										lowercase :							Dict    		=SequentialHook(PostForwardHook()		,					PostForwardHook()   )
										add_hook_to_module(__A		,					__A   )

										lowercase :							Union[str, Any]    		=test_model(__A   )
										assert torch.allclose(__A		,					output + 2		,					atol=1E-5   )







				def  lowerCamelCase_	(       self     :		Tuple   ):
										'''simple docstring'''
										lowercase :							Optional[Any]    		=ModelForTest()
										lowercase :							List[Any]    		=torch.randn(2		,					3   )
										lowercase :							int    		=test_model(__A   )

										lowercase :							Tuple    		=PostForwardHook()
										add_hook_to_module(__A		,					__A   )
										lowercase :							List[Any]    		=test_model(__A   )
										self.assertTrue(torch.allclose(__A		,					output + 1   )   )
										self.assertTrue(outputa.requires_grad   )

										lowercase :							int    		=True
										lowercase :							Any    		=test_model(__A   )
										self.assertFalse(outputa.requires_grad   )







				@require_multi_gpu
				def  lowerCamelCase_	(       self     :		str   ):
										'''simple docstring'''
										lowercase :							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
										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 :							Tuple    		=torch.randn(2		,					3   )
										lowercase :							Tuple    		=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   )   )
										lowercase :							List[str]    		=torch.randn(2		,					3   ).to(0   )
										lowercase :							Tuple    		=model(__A   )
										self.assertEqual(output.device		,					torch.device(0   )   )







				def  lowerCamelCase_	(       self     :		Any   ):
										'''simple docstring'''
										lowercase :							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
										lowercase :							Optional[int]    		={"""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
										lowercase :							List[str]    		=torch.device(hook_kwargs['''execution_device''']   )
										self.assertEqual(model.batchnorm.running_mean.device		,					__A   )

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

										# Now test with buffers included in the offload
										lowercase :							str    		={
										    """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'''   )   )

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







				def  lowerCamelCase_	(       self     :		List[str]   ):
										'''simple docstring'''
										lowercase :							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
										lowercase :							str    		=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
										lowercase :							List[Any]    		=torch.device(__A   )
										self.assertEqual(model.batchnorm.running_mean.device		,					__A   )

										lowercase :							List[str]    		=torch.randn(2		,					3   )
										lowercase :							Dict    		=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'''   )   )

										lowercase :							Dict    		=torch.randn(2		,					3   )
										lowercase :							List[str]    		=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  lowerCamelCase_	(       self     :		Optional[int]   ):
										'''simple docstring'''
										lowercase :							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
										lowercase :							int    		=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
										lowercase :							Optional[int]    		=torch.device(__A   )
										self.assertEqual(model.batchnorm.running_mean.device		,					__A   )

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

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








UpperCAmelCase_							:					dict[str, float]		 =   {
    "joule": 1.0,
    "kilojoule": 1_000,
    "megajoule": 1_000_000,
    "gigajoule": 1_000_000_000,
    "wattsecond": 1.0,
    "watthour": 3_600,
    "kilowatthour": 3_600_000,
    "newtonmeter": 1.0,
    "calorie_nutr": 4_186.8,
    "kilocalorie_nutr": 4_186_800.00,
    "electronvolt": 1.6_0217_6634e-19,
    "britishthermalunit_it": 1_055.05_585,
    "footpound": 1.35_5818,
}
def 							__SCREAMING_SNAKE_CASE     (   a__    :   str			,a__    :   str			,a__    :   float       )					->				float:
	if to_type not in ENERGY_CONVERSION or from_type not in ENERGY_CONVERSION:
		__A :			Optional[int]		    =			(
		    f"""Incorrect 'from_type' or 'to_type' value: {from_type!r}, {to_type!r}\n"""
		    f"""Valid values are: {", ".join(a__       )}"""
		)
		raise ValueError(a__       )
	return value * ENERGY_CONVERSION[from_type] / ENERGY_CONVERSION[to_type]


if __name__ == "__main__":
			import doctest

			doctest.testmod()





