






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



	'''simple docstring'''




	def __init__(			self   )			->       int:
								super().__init__()
								snake_case_           =				nn.Linear(3			,     4   )
								snake_case_           =				nn.BatchNormad(4   )
								snake_case_           =				nn.Linear(4			,     5   )






	def 				_UpperCamelCase     (			self			,     a   )			->       Dict:
								return self.lineara(self.batchnorm(self.lineara(lowercase_   )   )   )






class 					UpperCamelCase_ (     a_							):



	'''simple docstring'''




	def 				_UpperCamelCase     (			self			,     a			,     *a			,     **a   )			->       Union[str, Any]:
								return (args[0] + 1,) + args[1:], kwargs






class 					UpperCamelCase_ (     a_							):



	'''simple docstring'''




	def 				_UpperCamelCase     (			self			,     a			,     a   )			->       Tuple:
								return output + 1






class 					UpperCamelCase_ (     unittest.TestCase							):



	'''simple docstring'''




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

								snake_case_           =				PreForwardHook()
								add_hook_to_module(lowercase_			,     lowercase_   )
								snake_case_           =				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
								snake_case_           =				PreForwardHook()
								add_hook_to_module(lowercase_			,     lowercase_   )
								snake_case_           =				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
								snake_case_           =				SequentialHook(PreForwardHook()			,     PreForwardHook()   )
								add_hook_to_module(lowercase_			,     lowercase_   )

								snake_case_           =				test_model(lowercase_   )
								assert torch.allclose(lowercase_			,     lowercase_			,     atol=1E-5   )




	def 				_UpperCamelCase     (			self   )			->       Tuple:
								snake_case_           =				ModelForTest()
								snake_case_           =				torch.randn(2			,     3   )
								snake_case_           =				test_model(lowercase_   )

								snake_case_           =				PostForwardHook()
								add_hook_to_module(lowercase_			,     lowercase_   )
								snake_case_           =				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
								snake_case_           =				PostForwardHook()
								add_hook_to_module(lowercase_			,     lowercase_   )
								snake_case_           =				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
								snake_case_           =				SequentialHook(PostForwardHook()			,     PostForwardHook()   )
								add_hook_to_module(lowercase_			,     lowercase_   )

								snake_case_           =				test_model(lowercase_   )
								assert torch.allclose(lowercase_			,     output + 2			,     atol=1E-5   )




	def 				_UpperCamelCase     (			self   )			->       List[str]:
								snake_case_           =				ModelForTest()
								snake_case_           =				torch.randn(2			,     3   )
								snake_case_           =				test_model(lowercase_   )

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

								snake_case_           =				True
								snake_case_           =				test_model(lowercase_   )
								self.assertFalse(outputa.requires_grad   )




	@require_multi_gpu
	def 				_UpperCamelCase     (			self   )			->       List[str]:
								snake_case_           =				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
								snake_case_           =				torch.randn(2			,     3   )
								snake_case_           =				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_   )   )
								snake_case_           =				torch.randn(2			,     3   ).to(0   )
								snake_case_           =				model(lowercase_   )
								self.assertEqual(output.device			,     torch.device(0   )   )




	def 				_UpperCamelCase     (			self   )			->       Union[str, Any]:
								snake_case_           =				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
								snake_case_           =				{'''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
								snake_case_           =				torch.device(hook_kwargs['execution_device']   )
								self.assertEqual(model.batchnorm.running_mean.device			,     lowercase_   )

								snake_case_           =				torch.randn(2			,     3   )
								snake_case_           =				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
								snake_case_           =				{
								    '''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'   )   )

								snake_case_           =				torch.randn(2			,     3   )
								snake_case_           =				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 				_UpperCamelCase     (			self   )			->       Optional[Any]:
								snake_case_           =				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
								snake_case_           =				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
								snake_case_           =				torch.device(lowercase_   )
								self.assertEqual(model.batchnorm.running_mean.device			,     lowercase_   )

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

								snake_case_           =				torch.randn(2			,     3   )
								snake_case_           =				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 				_UpperCamelCase     (			self   )			->       str:
								snake_case_           =				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
								snake_case_           =				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
								snake_case_           =				torch.device(lowercase_   )
								self.assertEqual(model.batchnorm.running_mean.device			,     lowercase_   )

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

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








"""simple docstring"""





def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number < 0 or shift_amount < 0:
		raise ValueError('''both inputs must be positive integers'''					)

	lowerCAmelCase__      :		int    							=      str(bin(A_					)					)
	binary_number += "0" * shift_amount
	return binary_number

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number < 0 or shift_amount < 0:
		raise ValueError('''both inputs must be positive integers'''					)

	lowerCAmelCase__      :		Any    							=      str(bin(A_					)					)[2:]
	if shift_amount >= len(A_					):
		return "0b0"
	lowerCAmelCase__      :		Tuple    							=      binary_number[: len(A_					) - shift_amount]
	return "0b" + shifted_binary_number

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number >= 0:  # Get binary representation of positive number
		lowerCAmelCase__      :		Union[str, Any]    							=      '''0''' + str(bin(A_					)					).strip('''-'''					)[2:]
	else:  # Get binary (2's complement) representation of negative number
		lowerCAmelCase__      :		Optional[Any]    							=      len(bin(A_					)[3:]					)  # Find 2's complement of number
		lowerCAmelCase__      :		Optional[int]    							=      bin(abs(A_					) - (1 << binary_number_length)					)[3:]
		lowerCAmelCase__      :		int    							=      (
		    '''1''' + '''0''' * (binary_number_length - len(A_					)) + binary_number
		)

	if shift_amount >= len(A_					):
		return "0b" + binary_number[0] * len(A_					)
	return (
	    "0b"
	    + binary_number[0] * shift_amount
	    + binary_number[: len(A_					) - shift_amount]
	)


if __name__ == "__main__":
							import doctest

							doctest.testmod()




