


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




     '''simple docstring'''


     def __init__(						self)			->				Dict:
            super().__init__()
            _A      :       Optional[Any]				    = nn.Linear(3 ,						4)
            _A      :       Optional[int]				    = nn.BatchNormad(4)
            _A      :       int				    = nn.Linear(4 ,						5)


     def  _lowerCamelCase (						self ,						__lowerCamelCase)			->				str:
            return self.lineara(self.batchnorm(self.lineara(__lowerCamelCase)))







class      lowerCAmelCase__       (				a):




     '''simple docstring'''


     def  _lowerCamelCase (						self ,						__lowerCamelCase ,						*__lowerCamelCase ,						**__lowerCamelCase)			->				Optional[int]:
            return (args[0] + 1,) + args[1:], kwargs







class      lowerCAmelCase__       (				a):




     '''simple docstring'''


     def  _lowerCamelCase (						self ,						__lowerCamelCase ,						__lowerCamelCase)			->				Union[str, Any]:
            return output + 1







class      lowerCAmelCase__       (				unittest.TestCase):




     '''simple docstring'''


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

            _A      :       Optional[Any]				    = PreForwardHook()
            add_hook_to_module(__lowerCamelCase ,						__lowerCamelCase)
            _A      :       Optional[Any]				    = 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      :       Optional[Any]				    = PreForwardHook()
            add_hook_to_module(__lowerCamelCase ,						__lowerCamelCase)
            _A      :       Optional[int]				    = 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      :       Optional[Any]				    = SequentialHook(PreForwardHook() ,						PreForwardHook())
            add_hook_to_module(__lowerCamelCase ,						__lowerCamelCase)

            _A      :       List[str]				    = test_model(__lowerCamelCase)
            assert torch.allclose(__lowerCamelCase ,						__lowerCamelCase ,						atol=1e-5)


     def  _lowerCamelCase (						self)			->				Tuple:
            _A      :       Optional[int]				    = ModelForTest()
            _A      :       str				    = torch.randn(2 ,						3)
            _A      :       int				    = test_model(__lowerCamelCase)

            _A      :       Dict				    = PostForwardHook()
            add_hook_to_module(__lowerCamelCase ,						__lowerCamelCase)
            _A      :       Any				    = 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      :       List[str]				    = 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      :       Optional[int]				    = SequentialHook(PostForwardHook() ,						PostForwardHook())
            add_hook_to_module(__lowerCamelCase ,						__lowerCamelCase)

            _A      :       int				    = test_model(__lowerCamelCase)
            assert torch.allclose(__lowerCamelCase ,						output + 2 ,						atol=1e-5)


     def  _lowerCamelCase (						self)			->				Union[str, Any]:
            _A      :       int				    = ModelForTest()
            _A      :       str				    = torch.randn(2 ,						3)
            _A      :       Optional[Any]				    = test_model(__lowerCamelCase)

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

            _A      :       List[str]				    = True
            _A      :       Optional[int]				    = test_model(__lowerCamelCase)
            self.assertFalse(outputa.requires_grad)


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


     def  _lowerCamelCase (						self)			->				Optional[int]:
            _A      :       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
            _A      :       Union[str, Any]				    = {"""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      :       Any				    = torch.device(hook_kwargs["execution_device"])
            self.assertEqual(model.batchnorm.running_mean.device ,						__lowerCamelCase)

            _A      :       Union[str, Any]				    = 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_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      :       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(**__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      :       str				    = 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_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)			->				int:
            _A      :       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
            _A      :       Optional[int]				    = 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      :       Any				    = torch.device(__lowerCamelCase)
            self.assertEqual(model.batchnorm.running_mean.device ,						__lowerCamelCase)

            _A      :       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_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      :       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_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  _lowerCamelCase (						self)			->				str:
            _A      :       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
            _A      :       Any				    = 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      :       str				    = torch.device(__lowerCamelCase)
            self.assertEqual(model.batchnorm.running_mean.device ,						__lowerCamelCase)

            _A      :       Dict				    = torch.randn(2 ,						3)
            _A      :       Optional[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"))

            # 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      :       int				    = 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_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"""






import os
import sys
from contextlib import contextmanager


# Windows only
if os.name == "nt":
							import ctypes
							import msvcrt  # noqa





							class      lowerCamelCase__   (     ctypes.Structure				):







												"""simple docstring"""


												__a	       =							[("""size""", ctypes.c_int), ("""visible""", ctypes.c_byte)]
def      lowerCamelCase       (     )		->       Optional[int]:

							'''simple docstring'''



							if os.name == "nt":
														__UpperCAmelCase   :						Dict      							=					CursorInfo()
														__UpperCAmelCase   :						Any      							=					ctypes.windll.kernelaa.GetStdHandle(-1_1							)
														ctypes.windll.kernelaa.GetConsoleCursorInfo(_UpperCamelCase       ,			ctypes.byref(_UpperCamelCase							)							)
														__UpperCAmelCase   :						Tuple      							=					False
														ctypes.windll.kernelaa.SetConsoleCursorInfo(_UpperCamelCase       ,			ctypes.byref(_UpperCamelCase							)							)
							elif os.name == "posix":
														sys.stdout.write("""\033[?25l"""							)
														sys.stdout.flush()
def      lowerCamelCase       (     )		->       Optional[int]:

							'''simple docstring'''



							if os.name == "nt":
														__UpperCAmelCase   :						str      							=					CursorInfo()
														__UpperCAmelCase   :						int      							=					ctypes.windll.kernelaa.GetStdHandle(-1_1							)
														ctypes.windll.kernelaa.GetConsoleCursorInfo(_UpperCamelCase       ,			ctypes.byref(_UpperCamelCase							)							)
														__UpperCAmelCase   :						Union[str, Any]      							=					True
														ctypes.windll.kernelaa.SetConsoleCursorInfo(_UpperCamelCase       ,			ctypes.byref(_UpperCamelCase							)							)
							elif os.name == "posix":
														sys.stdout.write("""\033[?25h"""							)
														sys.stdout.flush()
@contextmanager
def      lowerCamelCase       (     )		->       str:

							'''simple docstring'''



							try:
														hide_cursor()
														yield
							finally:
														show_cursor()
