



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




    '''simple docstring'''






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






    def       _lowerCAmelCase (	self,		lowerCamelCase__		):
        return self.lineara(self.batchnorm(self.lineara(lowerCamelCase__		)		)		)






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__   ):




    '''simple docstring'''






    def       _lowerCAmelCase (	self,		lowerCamelCase__,		*lowerCamelCase__,		**lowerCamelCase__		):
        return (args[0] + 1,) + args[1:], kwargs






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__   ):




    '''simple docstring'''






    def       _lowerCAmelCase (	self,		lowerCamelCase__,		lowerCamelCase__		):
        return output + 1






class SCREAMING_SNAKE_CASE__					(		unittest.TestCase   ):




    '''simple docstring'''






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

        A			:	List[str]										=       PreForwardHook()
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)
        A			:	List[str]										=       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			:	Tuple										=       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			:	Any										=       SequentialHook(PreForwardHook(),		PreForwardHook()		)
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)

        A			:	Optional[Any]										=       test_model(lowerCamelCase__		)
        assert torch.allclose(lowerCamelCase__,		lowerCamelCase__,		atol=1e-5		)






    def       _lowerCAmelCase (	self		):
        A			:	Tuple										=       ModelForTest()
        A			:	Any										=       torch.randn(2,		3		)
        A			:	Any										=       test_model(lowerCamelCase__		)

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

        A			:	List[str]										=       test_model(lowerCamelCase__		)
        assert torch.allclose(lowerCamelCase__,		output + 2,		atol=1e-5		)






    def       _lowerCAmelCase (	self		):
        A			:	List[Any]										=       ModelForTest()
        A			:	Tuple										=       torch.randn(2,		3		)
        A			:	Union[str, Any]										=       test_model(lowerCamelCase__		)

        A			:	List[Any]										=       PostForwardHook()
        add_hook_to_module(lowerCamelCase__,		lowerCamelCase__		)
        A			:	List[Any]										=       test_model(lowerCamelCase__		)
        self.assertTrue(torch.allclose(lowerCamelCase__,		output + 1		)		)
        self.assertTrue(outputa.requires_grad		)

        A			:	int										=       True
        A			:	Tuple										=       test_model(lowerCamelCase__		)
        self.assertFalse(outputa.requires_grad		)






    @require_multi_gpu
    def       _lowerCAmelCase (	self		):
        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
        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			:	str										=       torch.randn(2,		3		)
        A			:	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__		)		)
        A			:	int										=       torch.randn(2,		3		).to(0		)
        A			:	str										=       model(lowerCamelCase__		)
        self.assertEqual(output.device,		torch.device(0		)		)






    def       _lowerCAmelCase (	self		):
        A			:	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
        A			:	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
        A			:	Dict										=       torch.device(hook_kwargs["""execution_device"""]		)
        self.assertEqual(model.batchnorm.running_mean.device,		lowerCamelCase__		)

        A			:	int										=       torch.randn(2,		3		)
        A			:	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
        A			:	int										=       {
            """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			:	int										=       torch.randn(2,		3		)
        A			:	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       _lowerCAmelCase (	self		):
        A			:	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
        A			:	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
        A			:	Optional[int]										=       torch.device(lowerCamelCase__		)
        self.assertEqual(model.batchnorm.running_mean.device,		lowerCamelCase__		)

        A			:	List[str]										=       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_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			:	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"""		)		)






    def       _lowerCAmelCase (	self		):
        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			:	List[str]										=       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			:	Any										=       torch.device(lowerCamelCase__		)
        self.assertEqual(model.batchnorm.running_mean.device,		lowerCamelCase__		)

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







import copy
from collections import OrderedDict
from typing import Mapping

from packaging import version

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


SCREAMING_SNAKE_CASE_:Optional[Any]					=						logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_:Union[str, Any]					=						{
    """microsoft/conditional-detr-resnet-50""": (
        """https://huggingface.co/microsoft/conditional-detr-resnet-50/resolve/main/config.json"""
    ),
}






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__   ):




    '''simple docstring'''






    __lowerCamelCase      :			Dict =				"conditional_detr"
    __lowerCamelCase      :			str =				["past_key_values"]
    __lowerCamelCase      :			str =				{
        "hidden_size": "d_model",
        "num_attention_heads": "encoder_attention_heads",
    }






    def __init__(	self,		lowerCamelCase__=True,		lowerCamelCase__=None,		lowerCamelCase__=3,		lowerCamelCase__=300,		lowerCamelCase__=6,		lowerCamelCase__=2048,		lowerCamelCase__=8,		lowerCamelCase__=6,		lowerCamelCase__=2048,		lowerCamelCase__=8,		lowerCamelCase__=0.0,		lowerCamelCase__=0.0,		lowerCamelCase__=True,		lowerCamelCase__="relu",		lowerCamelCase__=256,		lowerCamelCase__=0.1,		lowerCamelCase__=0.0,		lowerCamelCase__=0.0,		lowerCamelCase__=0.02,		lowerCamelCase__=1.0,		lowerCamelCase__=False,		lowerCamelCase__="sine",		lowerCamelCase__="resnet50",		lowerCamelCase__=True,		lowerCamelCase__=False,		lowerCamelCase__=2,		lowerCamelCase__=5,		lowerCamelCase__=2,		lowerCamelCase__=1,		lowerCamelCase__=1,		lowerCamelCase__=2,		lowerCamelCase__=5,		lowerCamelCase__=2,		lowerCamelCase__=0.25,		**lowerCamelCase__,		):
        if backbone_config is not None and use_timm_backbone:
            raise ValueError("""You can't specify both `backbone_config` and `use_timm_backbone`."""		)

        if not use_timm_backbone:
            if backbone_config is None:
                logger.info("""`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone."""		)
                A			:	List[Any]										=       CONFIG_MAPPING["""resnet"""](out_features=["""stage4"""]		)
            elif isinstance(lowerCamelCase__,		lowerCamelCase__		):
                A			:	Any										=       backbone_config.get("""model_type"""		)
                A			:	Optional[Any]										=       CONFIG_MAPPING[backbone_model_type]
                A			:	Tuple										=       config_class.from_dict(lowerCamelCase__		)

        A			:	Dict										=       use_timm_backbone
        A			:	int										=       backbone_config
        A			:	Union[str, Any]										=       num_channels
        A			:	Optional[Any]										=       num_queries
        A			:	Union[str, Any]										=       d_model
        A			:	str										=       encoder_ffn_dim
        A			:	List[Any]										=       encoder_layers
        A			:	Tuple										=       encoder_attention_heads
        A			:	Union[str, Any]										=       decoder_ffn_dim
        A			:	Tuple										=       decoder_layers
        A			:	int										=       decoder_attention_heads
        A			:	Union[str, Any]										=       dropout
        A			:	List[str]										=       attention_dropout
        A			:	Optional[int]										=       activation_dropout
        A			:	Optional[Any]										=       activation_function
        A			:	Any										=       init_std
        A			:	List[Any]										=       init_xavier_std
        A			:	Any										=       encoder_layerdrop
        A			:	List[str]										=       decoder_layerdrop
        A			:	int										=       encoder_layers
        A			:	Union[str, Any]										=       auxiliary_loss
        A			:	Union[str, Any]										=       position_embedding_type
        A			:	Tuple										=       backbone
        A			:	Dict										=       use_pretrained_backbone
        A			:	int										=       dilation
        # Hungarian matcher
        A			:	List[Any]										=       class_cost
        A			:	List[Any]										=       bbox_cost
        A			:	int										=       giou_cost
        # Loss coefficients
        A			:	List[Any]										=       mask_loss_coefficient
        A			:	Any										=       dice_loss_coefficient
        A			:	int										=       cls_loss_coefficient
        A			:	Tuple										=       bbox_loss_coefficient
        A			:	List[Any]										=       giou_loss_coefficient
        A			:	int										=       focal_alpha
        super().__init__(is_encoder_decoder=lowerCamelCase__,		**lowerCamelCase__		)






    @property
    def       _lowerCAmelCase (	self		):
        return self.encoder_attention_heads






    @property
    def       _lowerCAmelCase (	self		):
        return self.d_model






    def       _lowerCAmelCase (	self		):
        A			:	Dict										=       copy.deepcopy(self.__dict__		)
        if self.backbone_config is not None:
            A			:	List[Any]										=       self.backbone_config.to_dict()
        A			:	List[str]										=       self.__class__.model_type
        return output






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__   ):




    '''simple docstring'''






    __lowerCamelCase      :			Tuple =				version.parse("1.11"   )






    @property
    def       _lowerCAmelCase (	self		):
        return OrderedDict(
            [
                ("""pixel_values""", {0: """batch""", 1: """num_channels""", 2: """height""", 3: """width"""}),
                ("""pixel_mask""", {0: """batch"""}),
            ]		)






    @property
    def       _lowerCAmelCase (	self		):
        return 1e-5






    @property
    def       _lowerCAmelCase (	self		):
        return 12


