




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






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






							def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]       ,					_lowerCAmelCase:					List[Any]				):
									return self.lineara(self.batchnorm(self.lineara(__lowerCamelCase				)				)				)







class    __lowerCAmelCase       (							_A):






							def 					SCREAMING_SNAKE_CASE    (   self:					Tuple       ,					_lowerCAmelCase:					List[str]       ,					*_lowerCAmelCase:					str       ,					**_lowerCAmelCase:					Optional[Any]				):
									return (args[0] + 1,) + args[1:], kwargs







class    __lowerCAmelCase       (							_A):






							def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]       ,					_lowerCAmelCase:					Any       ,					_lowerCAmelCase:					Tuple				):
									return output + 1







class    __lowerCAmelCase       (							unittest.TestCase):






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

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

									lowercase    :Union[str, Any]      	=						test_model(__lowerCamelCase				)
									assert torch.allclose(__lowerCamelCase       ,					__lowerCamelCase       ,					atol=1e-5				)






							def 					SCREAMING_SNAKE_CASE    (   self:					Dict				):
									lowercase    :List[str]      	=						ModelForTest()
									lowercase    :Union[str, Any]      	=						torch.randn(2       ,					3				)
									lowercase    :Optional[int]      	=						test_model(__lowerCamelCase				)

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

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






							def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
									lowercase    :Tuple      	=						ModelForTest()
									lowercase    :List[str]      	=						torch.randn(2       ,					3				)
									lowercase    :int      	=						test_model(__lowerCamelCase				)

									lowercase    :Any      	=						PostForwardHook()
									add_hook_to_module(__lowerCamelCase       ,					__lowerCamelCase				)
									lowercase    :int      	=						test_model(__lowerCamelCase				)
									self.assertTrue(torch.allclose(__lowerCamelCase       ,					output + 1				)				)
									self.assertTrue(outputa.requires_grad				)

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






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






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

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

									lowercase    :int      	=						torch.randn(2       ,					3				)
									lowercase    :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 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
									lowercase    :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
									lowercase    :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
									lowercase    :Optional[int]      	=						torch.device(__lowerCamelCase				)
									self.assertEqual(model.batchnorm.running_mean.device       ,					__lowerCamelCase				)

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

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






							def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
									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
									lowercase    :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
									lowercase    :Dict      	=						torch.device(__lowerCamelCase				)
									self.assertEqual(model.batchnorm.running_mean.device       ,					__lowerCamelCase				)

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

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









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_UpperCAmelCase	:						List[Any]			=     {
    "configuration_swinv2": ["SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP", "Swinv2Config"],
}


try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCAmelCase	:						Union[str, Any]			=     [
         "SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST",
         "Swinv2ForImageClassification",
         "Swinv2ForMaskedImageModeling",
         "Swinv2Model",
         "Swinv2PreTrainedModel",
     ]


if TYPE_CHECKING:
     from .configuration_swinva import SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinvaConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_swinva import (
              SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST,
              SwinvaForImageClassification,
              SwinvaForMaskedImageModeling,
              SwinvaModel,
              SwinvaPreTrainedModel,
          )


else:
     import sys

     _UpperCAmelCase	:						Tuple			=     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



