





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






									"""simple docstring"""
									super().__init__()
									lowercase_	:   Union[str, Any]  				=     nn.Linear(3,   4  )
									lowercase_	:   Any  				=     nn.BatchNormad(4  )
									lowercase_	:   Union[str, Any]  				=     nn.Linear(4,   5  )






		def 						snake_case__				(							self,   snake_case__  )       -> int:






									"""simple docstring"""
									return self.lineara(self.batchnorm(self.lineara(__lowerCAmelCase  )  )  )




class     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






		def 						snake_case__				(							self,   snake_case__,   *snake_case__,   **snake_case__  )       -> Any:






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




class     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






		def 						snake_case__				(							self,   snake_case__,   snake_case__  )       -> Tuple:






									"""simple docstring"""
									return output + 1




class     UpperCamelCase__					(     unittest.TestCase				):






		'''simple docstring'''






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   List[str]  				=     ModelForTest()
									lowercase_	:   Tuple  				=     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 						snake_case__				(							self  )       -> str:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     ModelForTest()
									lowercase_	:   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 						snake_case__				(							self  )       -> List[str]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     ModelForTest()
									lowercase_	:   Union[str, Any]  				=     torch.randn(2,   3  )
									lowercase_	:   int  				=     test_model(x + 1  )
									lowercase_	:   List[Any]  				=     test_model(x + 2  )

									lowercase_	:   Optional[int]  				=     PreForwardHook()
									add_hook_to_module(__lowerCAmelCase,   __lowerCAmelCase  )
									lowercase_	:   Dict  				=     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_	:   Tuple  				=     PreForwardHook()
									add_hook_to_module(__lowerCAmelCase,   __lowerCAmelCase  )
									lowercase_	:   List[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_	:   Dict  				=     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 						snake_case__				(							self  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   Dict  				=     ModelForTest()
									lowercase_	:   int  				=     torch.randn(2,   3  )
									lowercase_	:   Union[str, Any]  				=     test_model(__lowerCAmelCase  )

									lowercase_	:   List[str]  				=     PostForwardHook()
									add_hook_to_module(__lowerCAmelCase,   __lowerCAmelCase  )
									lowercase_	:   List[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
									lowercase_	:   int  				=     PostForwardHook()
									add_hook_to_module(__lowerCAmelCase,   __lowerCAmelCase  )
									lowercase_	:   List[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_	:   int  				=     SequentialHook(PostForwardHook(),   PostForwardHook()  )
									add_hook_to_module(__lowerCAmelCase,   __lowerCAmelCase  )

									lowercase_	:   Optional[Any]  				=     test_model(__lowerCAmelCase  )
									assert torch.allclose(__lowerCAmelCase,   output + 2,   atol=1E-5  )






		def 						snake_case__				(							self  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     ModelForTest()
									lowercase_	:   Optional[Any]  				=     torch.randn(2,   3  )
									lowercase_	:   List[Any]  				=     test_model(__lowerCAmelCase  )

									lowercase_	:   Optional[int]  				=     PostForwardHook()
									add_hook_to_module(__lowerCAmelCase,   __lowerCAmelCase  )
									lowercase_	:   Tuple  				=     test_model(__lowerCAmelCase  )
									self.assertTrue(torch.allclose(__lowerCAmelCase,   output + 1  )  )
									self.assertTrue(outputa.requires_grad  )

									lowercase_	:   Tuple  				=     True
									lowercase_	:   Optional[int]  				=     test_model(__lowerCAmelCase  )
									self.assertFalse(outputa.requires_grad  )






		@require_multi_gpu
		def 						snake_case__				(							self  )       -> Optional[Any]:






									"""simple docstring"""
									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
									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_	:   List[str]  				=     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_	:   Tuple  				=     torch.randn(2,   3  ).to(0  )
									lowercase_	:   int  				=     model(__lowerCAmelCase  )
									self.assertEqual(output.device,   torch.device(0  )  )






		def 						snake_case__				(							self  )       -> Tuple:






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

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

									# Now test with buffers included in the offload
									lowercase_	:   Optional[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_	:   List[Any]  				=     torch.randn(2,   3  )
									lowercase_	:   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 						snake_case__				(							self  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   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
									lowercase_	:   Union[str, Any]  				=     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[Any]  				=     torch.device(__lowerCAmelCase  )
									self.assertEqual(model.batchnorm.running_mean.device,   __lowerCAmelCase  )

									lowercase_	:   List[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"""  )  )

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






		def 						snake_case__				(							self  )       -> Tuple:






									"""simple docstring"""
									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_	:   Optional[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_	:   Union[str, Any]  				=     torch.device(__lowerCAmelCase  )
									self.assertEqual(model.batchnorm.running_mean.device,   __lowerCAmelCase  )

									lowercase_	:   int  				=     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"""  )  )

									# 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_	:   List[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"""  )  )
"""simple docstring"""






import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


__lowerCAmelCase       :			Tuple   	=logging.get_logger(__name__)

__lowerCAmelCase       :			Dict   	={
    """ut/deta""": """https://huggingface.co/ut/deta/resolve/main/config.json""",
}





class 					_A		(			lowerCAmelCase      ):
      snake_case__     :  Optional[int]				    =      'deta'
      snake_case__     :  Dict				    =      {
          'hidden_size': 'd_model',
          'num_attention_heads': 'encoder_attention_heads',
      }

      def __init__(		self						,			__lowerCAmelCase=None						,			__lowerCAmelCase=900						,			__lowerCAmelCase=2048						,			__lowerCAmelCase=6						,			__lowerCAmelCase=2048						,			__lowerCAmelCase=8						,			__lowerCAmelCase=6						,			__lowerCAmelCase=1024						,			__lowerCAmelCase=8						,			__lowerCAmelCase=0.0						,			__lowerCAmelCase=True						,			__lowerCAmelCase="relu"						,			__lowerCAmelCase=256						,			__lowerCAmelCase=0.1						,			__lowerCAmelCase=0.0						,			__lowerCAmelCase=0.0						,			__lowerCAmelCase=0.0_2						,			__lowerCAmelCase=1.0						,			__lowerCAmelCase=True						,			__lowerCAmelCase=False						,			__lowerCAmelCase="sine"						,			__lowerCAmelCase=5						,			__lowerCAmelCase=4						,			__lowerCAmelCase=4						,			__lowerCAmelCase=True						,			__lowerCAmelCase=300						,			__lowerCAmelCase=True						,			__lowerCAmelCase=True						,			__lowerCAmelCase=1						,			__lowerCAmelCase=5						,			__lowerCAmelCase=2						,			__lowerCAmelCase=1						,			__lowerCAmelCase=1						,			__lowerCAmelCase=5						,			__lowerCAmelCase=2						,			__lowerCAmelCase=0.1						,			__lowerCAmelCase=0.2_5						,			**__lowerCAmelCase						,			):







            """simple docstring"""



            if backbone_config is None:
                  logger.info("""`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone."""		)
                  lowercase								=      CONFIG_MAPPING["""resnet"""](out_features=["""stage2""", """stage3""", """stage4"""]		)
            else:
                  if isinstance(__lowerCAmelCase						,			__lowerCAmelCase		):
                        lowercase								=      backbone_config.pop("""model_type"""		)
                        lowercase								=      CONFIG_MAPPING[backbone_model_type]
                        lowercase								=      config_class.from_dict(__lowerCAmelCase		)

            lowercase								=      backbone_config
            lowercase								=      num_queries
            lowercase								=      max_position_embeddings
            lowercase								=      d_model
            lowercase								=      encoder_ffn_dim
            lowercase								=      encoder_layers
            lowercase								=      encoder_attention_heads
            lowercase								=      decoder_ffn_dim
            lowercase								=      decoder_layers
            lowercase								=      decoder_attention_heads
            lowercase								=      dropout
            lowercase								=      attention_dropout
            lowercase								=      activation_dropout
            lowercase								=      activation_function
            lowercase								=      init_std
            lowercase								=      init_xavier_std
            lowercase								=      encoder_layerdrop
            lowercase								=      auxiliary_loss
            lowercase								=      position_embedding_type
            # deformable attributes
            lowercase								=      num_feature_levels
            lowercase								=      encoder_n_points
            lowercase								=      decoder_n_points
            lowercase								=      two_stage
            lowercase								=      two_stage_num_proposals
            lowercase								=      with_box_refine
            lowercase								=      assign_first_stage
            if two_stage is True and with_box_refine is False:
                  raise ValueError("""If two_stage is True, with_box_refine must be True."""		)
            # Hungarian matcher
            lowercase								=      class_cost
            lowercase								=      bbox_cost
            lowercase								=      giou_cost
            # Loss coefficients
            lowercase								=      mask_loss_coefficient
            lowercase								=      dice_loss_coefficient
            lowercase								=      bbox_loss_coefficient
            lowercase								=      giou_loss_coefficient
            lowercase								=      eos_coefficient
            lowercase								=      focal_alpha
            super().__init__(is_encoder_decoder=__lowerCAmelCase						,			**__lowerCAmelCase		)

      @property
      def 							A__			(		self		):







            """simple docstring"""



            return self.encoder_attention_heads

      @property
      def 							A__			(		self		):







            """simple docstring"""



            return self.d_model







      def 							A__			(		self		):







            """simple docstring"""



            lowercase								=      copy.deepcopy(self.__dict__		)
            lowercase								=      self.backbone_config.to_dict()
            lowercase								=      self.__class__.model_type
            return output



