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





						def __init__(self						)				->				Any:
												super().__init__()
												UpperCamelCase_		:							int  =		nn.Linear(3			,     4						)
												UpperCamelCase_		:							Union[str, Any]  =		nn.BatchNormad(4						)
												UpperCamelCase_		:							List[str]  =		nn.Linear(4			,     5						)





						def        A_							(self			,     __UpperCamelCase						)				->				Optional[int]:
												return self.lineara(self.batchnorm(self.lineara(__SCREAMING_SNAKE_CASE						)						)						)


class        UpperCamelCase						(				_lowercase				):





						def        A_							(self			,     __UpperCamelCase			,     *__UpperCamelCase			,     **__UpperCamelCase						)				->				str:
												return (args[0] + 1,) + args[1:], kwargs


class        UpperCamelCase						(				_lowercase				):





						def        A_							(self			,     __UpperCamelCase			,     __UpperCamelCase						)				->				int:
												return output + 1


class        UpperCamelCase						(				unittest.TestCase				):





						def        A_							(self						)				->				List[str]:
												UpperCamelCase_		:							Dict  =		ModelForTest()
												UpperCamelCase_		:							Dict  =		ModelHook()

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





						def        A_							(self						)				->				Optional[Any]:
												UpperCamelCase_		:							List[str]  =		ModelForTest()
												UpperCamelCase_		:							List[str]  =		ModelHook()

												add_hook_to_module(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE						)
												add_hook_to_module(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     append=__SCREAMING_SNAKE_CASE						)

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





						def        A_							(self						)				->				Any:
												UpperCamelCase_		:							Optional[int]  =		ModelForTest()
												UpperCamelCase_		:							Optional[Any]  =		torch.randn(2			,     3						)
												UpperCamelCase_		:							List[Any]  =		test_model(x + 1						)
												UpperCamelCase_		:							Dict  =		test_model(x + 2						)

												UpperCamelCase_		:							Dict  =		PreForwardHook()
												add_hook_to_module(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE						)
												UpperCamelCase_		:							Any  =		test_model(__SCREAMING_SNAKE_CASE						)
												self.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     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(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE						)
												UpperCamelCase_		:							str  =		test_model(__SCREAMING_SNAKE_CASE						)
												self.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     atol=1E-5						)						)

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

												UpperCamelCase_		:							List[Any]  =		test_model(__SCREAMING_SNAKE_CASE						)
												assert torch.allclose(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     atol=1E-5						)





						def        A_							(self						)				->				Optional[Any]:
												UpperCamelCase_		:							Optional[int]  =		ModelForTest()
												UpperCamelCase_		:							Optional[int]  =		torch.randn(2			,     3						)
												UpperCamelCase_		:							List[Any]  =		test_model(__SCREAMING_SNAKE_CASE						)

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

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

												# You need to use the sequential hook to chain two or more hooks
												UpperCamelCase_		:							Optional[int]  =		SequentialHook(PostForwardHook()			,     PostForwardHook()						)
												add_hook_to_module(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE						)

												UpperCamelCase_		:							Tuple  =		test_model(__SCREAMING_SNAKE_CASE						)
												assert torch.allclose(__SCREAMING_SNAKE_CASE			,     output + 2			,     atol=1E-5						)





						def        A_							(self						)				->				List[str]:
												UpperCamelCase_		:							Optional[int]  =		ModelForTest()
												UpperCamelCase_		:							Optional[Any]  =		torch.randn(2			,     3						)
												UpperCamelCase_		:							Any  =		test_model(__SCREAMING_SNAKE_CASE						)

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

												UpperCamelCase_		:							Dict  =		True
												UpperCamelCase_		:							Optional[Any]  =		test_model(__SCREAMING_SNAKE_CASE						)
												self.assertFalse(outputa.requires_grad						)





						@require_multi_gpu
						def        A_							(self						)				->				Dict:
												UpperCamelCase_		:							List[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_		:							Dict  =		torch.randn(2			,     3						)
												UpperCamelCase_		:							Dict  =		model(__SCREAMING_SNAKE_CASE						)
												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(__SCREAMING_SNAKE_CASE			,     AlignDevicesHook(io_same_device=__SCREAMING_SNAKE_CASE						)						)
												UpperCamelCase_		:							str  =		torch.randn(2			,     3						).to(0						)
												UpperCamelCase_		:							Optional[Any]  =		model(__SCREAMING_SNAKE_CASE						)
												self.assertEqual(output.device			,     torch.device(0						)						)





						def        A_							(self						)				->				Any:
												UpperCamelCase_		:							List[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
												UpperCamelCase_		:							Union[str, Any]  =		{"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

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

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

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

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

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

												# 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_		:							Tuple  =		torch.randn(2			,     3						)
												UpperCamelCase_		:							Tuple  =		model(__SCREAMING_SNAKE_CASE						)
												self.assertEqual(output.device			,     __SCREAMING_SNAKE_CASE						)

												# 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						)				->				Union[str, Any]:
												UpperCamelCase_		:							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
												UpperCamelCase_		:							Optional[Any]  =		0 if torch.cuda.is_available() else """cpu"""
												attach_align_device_hook(__SCREAMING_SNAKE_CASE			,     execution_device=__SCREAMING_SNAKE_CASE			,     offload=__SCREAMING_SNAKE_CASE						)

												# 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_		:							str  =		torch.device(__SCREAMING_SNAKE_CASE						)
												self.assertEqual(model.batchnorm.running_mean.device			,     __SCREAMING_SNAKE_CASE						)

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

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

												# 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[Any]  =		model(__SCREAMING_SNAKE_CASE						)
												self.assertEqual(output.device			,     __SCREAMING_SNAKE_CASE						)

												# Removing hooks loads back the weights in the model.
												remove_hook_from_submodules(__SCREAMING_SNAKE_CASE						)
												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						)				->				Any:
												UpperCamelCase_		:							Tuple  =		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(
												    __SCREAMING_SNAKE_CASE			,     execution_device=__SCREAMING_SNAKE_CASE			,     offload=__SCREAMING_SNAKE_CASE			,     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_		:							Tuple  =		torch.device(__SCREAMING_SNAKE_CASE						)
												self.assertEqual(model.batchnorm.running_mean.device			,     __SCREAMING_SNAKE_CASE						)

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

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

												# 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_		:							Optional[int]  =		torch.randn(2			,     3						)
												UpperCamelCase_		:							List[str]  =		model(__SCREAMING_SNAKE_CASE						)
												self.assertEqual(output.device			,     __SCREAMING_SNAKE_CASE						)

												# Removing hooks loads back the weights in the model.
												remove_hook_from_submodules(__SCREAMING_SNAKE_CASE						)
												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 math import ceil



def     __magic_name__		(					UpperCamelCase     :		int = 1001     )     ->						int:
    a__		=			1

    for i in range(1							,		int(ceil(n / 2.0     )     )     ):
        a__		=			2 * i + 1
        a__		=			2 * i
        a__		=			total + 4 * odd**2 - 6 * even

    return total


if __name__ == "__main__":
    import sys

    if len(sys.argv) == 1:
        print(solution())
    else:
        try:
            a		:       Any     						=       int(sys.argv[1])
            print(solution(n))
        except ValueError:
            print('Invalid entry - please enter a number')


