





"""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):







				def __init__(							self   :					Any       ):



					"""simple docstring"""

					super().__init__()
					A_  :				int	   =		nn.Linear(3      ,     4       )
					A_  :				str	   =		nn.BatchNormad(4       )
					A_  :				str	   =		nn.Linear(4      ,     5       )







				def 					a_  (							self   :					str      ,     _lowerCamelCase   :					Optional[Any]       ):



					"""simple docstring"""

					return self.lineara(self.batchnorm(self.lineara(_lowerCamelCase       )       )       )

class        lowercase      (    __UpperCAmelCase):







				def 					a_  (							self   :					Any      ,     _lowerCamelCase   :					Any      ,     *_lowerCamelCase   :					Optional[Any]      ,     **_lowerCamelCase   :					List[str]       ):



					"""simple docstring"""

					return (args[0] + 1,) + args[1:], kwargs

class        lowercase      (    __UpperCAmelCase):







				def 					a_  (							self   :					Tuple      ,     _lowerCamelCase   :					int      ,     _lowerCamelCase   :					Tuple       ):



					"""simple docstring"""

					return output + 1

class        lowercase      (    unittest.TestCase):







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

					A_  :				Optional[Any]	   =		ModelForTest()
					A_  :				Optional[int]	   =		ModelHook()

					add_hook_to_module(_lowerCamelCase      ,     _lowerCamelCase       )
					self.assertEqual(test_model._hf_hook      ,     _lowerCamelCase       )
					self.assertTrue(hasattr(_lowerCamelCase      ,     '''_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(_lowerCamelCase       )
					self.assertFalse(hasattr(_lowerCamelCase      ,     '''_hf_hook'''       )       )
					self.assertFalse(hasattr(_lowerCamelCase      ,     '''_old_forward'''       )       )







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

					A_  :				Optional[Any]	   =		ModelForTest()
					A_  :				Optional[int]	   =		ModelHook()

					add_hook_to_module(_lowerCamelCase      ,     _lowerCamelCase       )
					add_hook_to_module(_lowerCamelCase      ,     _lowerCamelCase      ,     append=_lowerCamelCase       )

					self.assertEqual(isinstance(test_model._hf_hook      ,     _lowerCamelCase       )      ,     _lowerCamelCase       )
					self.assertEqual(len(test_model._hf_hook.hooks       )      ,     2       )
					self.assertTrue(hasattr(_lowerCamelCase      ,     '''_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(_lowerCamelCase       )
					self.assertFalse(hasattr(_lowerCamelCase      ,     '''_hf_hook'''       )       )
					self.assertFalse(hasattr(_lowerCamelCase      ,     '''_old_forward'''       )       )







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

					A_  :				List[Any]	   =		ModelForTest()
					A_  :				List[Any]	   =		torch.randn(2      ,     3       )
					A_  :				Optional[Any]	   =		test_model(x + 1       )
					A_  :				List[Any]	   =		test_model(x + 2       )

					A_  :				Optional[int]	   =		PreForwardHook()
					add_hook_to_module(_lowerCamelCase      ,     _lowerCamelCase       )
					A_  :				Tuple	   =		test_model(_lowerCamelCase       )
					self.assertTrue(torch.allclose(_lowerCamelCase      ,     _lowerCamelCase      ,     atol=1E-5       )       )

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

					# You need to use the sequential hook to chain two or more hooks
					A_  :				int	   =		SequentialHook(PreForwardHook()      ,     PreForwardHook()       )
					add_hook_to_module(_lowerCamelCase      ,     _lowerCamelCase       )

					A_  :				Any	   =		test_model(_lowerCamelCase       )
					assert torch.allclose(_lowerCamelCase      ,     _lowerCamelCase      ,     atol=1E-5       )







				def 					a_  (							self   :					Union[str, Any]       ):



					"""simple docstring"""

					A_  :				int	   =		ModelForTest()
					A_  :				Any	   =		torch.randn(2      ,     3       )
					A_  :				Tuple	   =		test_model(_lowerCamelCase       )

					A_  :				Optional[int]	   =		PostForwardHook()
					add_hook_to_module(_lowerCamelCase      ,     _lowerCamelCase       )
					A_  :				int	   =		test_model(_lowerCamelCase       )
					self.assertTrue(torch.allclose(_lowerCamelCase      ,     output + 1      ,     atol=1E-5       )       )

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

					# You need to use the sequential hook to chain two or more hooks
					A_  :				Any	   =		SequentialHook(PostForwardHook()      ,     PostForwardHook()       )
					add_hook_to_module(_lowerCamelCase      ,     _lowerCamelCase       )

					A_  :				str	   =		test_model(_lowerCamelCase       )
					assert torch.allclose(_lowerCamelCase      ,     output + 2      ,     atol=1E-5       )







				def 					a_  (							self   :					Any       ):



					"""simple docstring"""

					A_  :				Optional[Any]	   =		ModelForTest()
					A_  :				List[Any]	   =		torch.randn(2      ,     3       )
					A_  :				Optional[Any]	   =		test_model(_lowerCamelCase       )

					A_  :				Union[str, Any]	   =		PostForwardHook()
					add_hook_to_module(_lowerCamelCase      ,     _lowerCamelCase       )
					A_  :				Optional[Any]	   =		test_model(_lowerCamelCase       )
					self.assertTrue(torch.allclose(_lowerCamelCase      ,     output + 1       )       )
					self.assertTrue(outputa.requires_grad       )

					A_  :				int	   =		True
					A_  :				Dict	   =		test_model(_lowerCamelCase       )
					self.assertFalse(outputa.requires_grad       )







				@require_multi_gpu
				def 					a_  (							self   :					int       ):



					"""simple docstring"""

					A_  :				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
					A_  :				int	   =		torch.randn(2      ,     3       )
					A_  :				int	   =		model(_lowerCamelCase       )
					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(_lowerCamelCase      ,     AlignDevicesHook(io_same_device=_lowerCamelCase       )       )
					A_  :				List[str]	   =		torch.randn(2      ,     3       ).to(0       )
					A_  :				Tuple	   =		model(_lowerCamelCase       )
					self.assertEqual(output.device      ,     torch.device(0       )       )







				def 					a_  (							self   :					Dict       ):



					"""simple docstring"""

					A_  :				Optional[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
					A_  :				Optional[int]	   =		{'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

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

					# 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
					A_  :				Tuple	   =		torch.device(hook_kwargs['''execution_device''']       )
					self.assertEqual(model.batchnorm.running_mean.device      ,     _lowerCamelCase       )

					A_  :				List[str]	   =		torch.randn(2      ,     3       )
					A_  :				Any	   =		model(_lowerCamelCase       )
					self.assertEqual(output.device      ,     _lowerCamelCase       )

					# 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
					A_  :				List[str]	   =		{
					    '''execution_device''': 0 if torch.cuda.is_available() else '''cpu''',
					    '''offload''': True,
					    '''offload_buffers''': True,
					}

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

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

					A_  :				List[str]	   =		torch.randn(2      ,     3       )
					A_  :				List[str]	   =		model(_lowerCamelCase       )
					self.assertEqual(output.device      ,     _lowerCamelCase       )

					# 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   :					List[str]       ):



					"""simple docstring"""

					A_  :				Optional[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
					A_  :				Tuple	   =		0 if torch.cuda.is_available() else '''cpu'''
					attach_align_device_hook(_lowerCamelCase      ,     execution_device=_lowerCamelCase      ,     offload=_lowerCamelCase       )

					# 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
					A_  :				Tuple	   =		torch.device(_lowerCamelCase       )
					self.assertEqual(model.batchnorm.running_mean.device      ,     _lowerCamelCase       )

					A_  :				Union[str, Any]	   =		torch.randn(2      ,     3       )
					A_  :				Dict	   =		model(_lowerCamelCase       )
					self.assertEqual(output.device      ,     _lowerCamelCase       )

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

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

					A_  :				str	   =		torch.randn(2      ,     3       )
					A_  :				str	   =		model(_lowerCamelCase       )
					self.assertEqual(output.device      ,     _lowerCamelCase       )

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(_lowerCamelCase       )
					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   :					Union[str, Any]       ):



					"""simple docstring"""

					A_  :				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
					A_  :				int	   =		0 if torch.cuda.is_available() else '''cpu'''
					attach_align_device_hook(
					    _lowerCamelCase      ,     execution_device=_lowerCamelCase      ,     offload=_lowerCamelCase      ,     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
					A_  :				Dict	   =		torch.device(_lowerCamelCase       )
					self.assertEqual(model.batchnorm.running_mean.device      ,     _lowerCamelCase       )

					A_  :				List[str]	   =		torch.randn(2      ,     3       )
					A_  :				Optional[int]	   =		model(_lowerCamelCase       )
					self.assertEqual(output.device      ,     _lowerCamelCase       )

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

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

					A_  :				List[Any]	   =		torch.randn(2      ,     3       )
					A_  :				Any	   =		model(_lowerCamelCase       )
					self.assertEqual(output.device      ,     _lowerCamelCase       )

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

from __future__ import annotations

from collections import namedtuple

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	A_  :				int	   =		namedtuple('''result'''  ,      '''name value'''	)
	if (voltage, current, power).count(0	) != 1:
		raise ValueError('''Only one argument must be 0'''	)
	elif power < 0:
		raise ValueError(
		    '''Power cannot be negative in any electrical/electronics system'''	)
	elif voltage == 0:
		return result('''voltage'''  ,      power / current	)
	elif current == 0:
		return result('''current'''  ,      power / voltage	)
	elif power == 0:
		return result('''power'''  ,      float(round(abs(voltage * current	)  ,      2	)	)	)
	else:
		raise ValueError('''Exactly one argument must be 0'''	)


if __name__ == "__main__":
		import doctest

		doctest.testmod()





