

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


							"""simple docstring"""





							def __init__( self							):
												super().__init__()
												__lowerCAmelCase					=							nn.Linear(3					,						4							)
												__lowerCAmelCase					=							nn.BatchNormad(4							)
												__lowerCAmelCase					=							nn.Linear(4					,						5							)




							def        A__  ( self					,						_A							):
												return self.lineara(self.batchnorm(self.lineara(lowerCAmelCase_							)							)							)





class     snake_case_			(				__a							):


							"""simple docstring"""





							def        A__  ( self					,						_A					,						*_A					,						**_A							):
												return (args[0] + 1,) + args[1:], kwargs





class     snake_case_			(				__a							):


							"""simple docstring"""





							def        A__  ( self					,						_A					,						_A							):
												return output + 1





class     snake_case_			(				unittest.TestCase							):


							"""simple docstring"""





							def        A__  ( self							):
												__lowerCAmelCase					=							ModelForTest()
												__lowerCAmelCase					=							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        A__  ( self							):
												__lowerCAmelCase					=							ModelForTest()
												__lowerCAmelCase					=							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        A__  ( self							):
												__lowerCAmelCase					=							ModelForTest()
												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							test_model(x + 1							)
												__lowerCAmelCase					=							test_model(x + 2							)

												__lowerCAmelCase					=							PreForwardHook()
												add_hook_to_module(lowerCAmelCase_					,						lowerCAmelCase_							)
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							PreForwardHook()
												add_hook_to_module(lowerCAmelCase_					,						lowerCAmelCase_							)
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							SequentialHook(PreForwardHook()					,						PreForwardHook()							)
												add_hook_to_module(lowerCAmelCase_					,						lowerCAmelCase_							)

												__lowerCAmelCase					=							test_model(lowerCAmelCase_							)
												assert torch.allclose(lowerCAmelCase_					,						lowerCAmelCase_					,						atol=1e-5							)




							def        A__  ( self							):
												__lowerCAmelCase					=							ModelForTest()
												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							test_model(lowerCAmelCase_							)

												__lowerCAmelCase					=							PostForwardHook()
												add_hook_to_module(lowerCAmelCase_					,						lowerCAmelCase_							)
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							PostForwardHook()
												add_hook_to_module(lowerCAmelCase_					,						lowerCAmelCase_							)
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							SequentialHook(PostForwardHook()					,						PostForwardHook()							)
												add_hook_to_module(lowerCAmelCase_					,						lowerCAmelCase_							)

												__lowerCAmelCase					=							test_model(lowerCAmelCase_							)
												assert torch.allclose(lowerCAmelCase_					,						output + 2					,						atol=1e-5							)




							def        A__  ( self							):
												__lowerCAmelCase					=							ModelForTest()
												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							test_model(lowerCAmelCase_							)

												__lowerCAmelCase					=							PostForwardHook()
												add_hook_to_module(lowerCAmelCase_					,						lowerCAmelCase_							)
												__lowerCAmelCase					=							test_model(lowerCAmelCase_							)
												self.assertTrue(torch.allclose(lowerCAmelCase_					,						output + 1							)							)
												self.assertTrue(outputa.requires_grad							)

												__lowerCAmelCase					=							True
												__lowerCAmelCase					=							test_model(lowerCAmelCase_							)
												self.assertFalse(outputa.requires_grad							)




							@require_multi_gpu
							def        A__  ( self							):
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							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_							)							)
												__lowerCAmelCase					=							torch.randn(2					,						3							).to(0							)
												__lowerCAmelCase					=							model(lowerCAmelCase_							)
												self.assertEqual(output.device					,						torch.device(0							)							)




							def        A__  ( self							):
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							{'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
												__lowerCAmelCase					=							torch.device(hook_kwargs['execution_device']							)
												self.assertEqual(model.batchnorm.running_mean.device					,						lowerCAmelCase_							)

												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							{
												    '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'							)							)

												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							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        A__  ( self							):
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							torch.device(lowerCAmelCase_							)
												self.assertEqual(model.batchnorm.running_mean.device					,						lowerCAmelCase_							)

												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							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'							)							)

												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							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        A__  ( self							):
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							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
												__lowerCAmelCase					=							torch.device(lowerCAmelCase_							)
												self.assertEqual(model.batchnorm.running_mean.device					,						lowerCAmelCase_							)

												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							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'							)							)

												__lowerCAmelCase					=							torch.randn(2					,						3							)
												__lowerCAmelCase					=							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_tf_available, is_torch_available


lowerCamelCase      				=     {
    '''configuration_data2vec_audio''': ['''DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''Data2VecAudioConfig'''],
    '''configuration_data2vec_text''': [
        '''DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''Data2VecTextConfig''',
        '''Data2VecTextOnnxConfig''',
    ],
    '''configuration_data2vec_vision''': [
        '''DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''Data2VecVisionConfig''',
        '''Data2VecVisionOnnxConfig''',
    ],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowerCamelCase      				=     [
		    '''DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''Data2VecAudioForAudioFrameClassification''',
		    '''Data2VecAudioForCTC''',
		    '''Data2VecAudioForSequenceClassification''',
		    '''Data2VecAudioForXVector''',
		    '''Data2VecAudioModel''',
		    '''Data2VecAudioPreTrainedModel''',
		]
		lowerCamelCase      				=     [
		    '''DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''Data2VecTextForCausalLM''',
		    '''Data2VecTextForMaskedLM''',
		    '''Data2VecTextForMultipleChoice''',
		    '''Data2VecTextForQuestionAnswering''',
		    '''Data2VecTextForSequenceClassification''',
		    '''Data2VecTextForTokenClassification''',
		    '''Data2VecTextModel''',
		    '''Data2VecTextPreTrainedModel''',
		]
		lowerCamelCase      				=     [
		    '''DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''Data2VecVisionForImageClassification''',
		    '''Data2VecVisionForMaskedImageModeling''',
		    '''Data2VecVisionForSemanticSegmentation''',
		    '''Data2VecVisionModel''',
		    '''Data2VecVisionPreTrainedModel''',
		]

if is_tf_available():
		lowerCamelCase      				=     [
		    '''TFData2VecVisionForImageClassification''',
		    '''TFData2VecVisionForSemanticSegmentation''',
		    '''TFData2VecVisionModel''',
		    '''TFData2VecVisionPreTrainedModel''',
		]

if TYPE_CHECKING:
		from .configuration_dataavec_audio import DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP, DataaVecAudioConfig
		from .configuration_dataavec_text import (
		    DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    DataaVecTextConfig,
		    DataaVecTextOnnxConfig,
		)
		from .configuration_dataavec_vision import (
		    DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    DataaVecVisionConfig,
		    DataaVecVisionOnnxConfig,
		)

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_dataavec_audio import (
				    DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DataaVecAudioForAudioFrameClassification,
				    DataaVecAudioForCTC,
				    DataaVecAudioForSequenceClassification,
				    DataaVecAudioForXVector,
				    DataaVecAudioModel,
				    DataaVecAudioPreTrainedModel,
				)
				from .modeling_dataavec_text import (
				    DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DataaVecTextForCausalLM,
				    DataaVecTextForMaskedLM,
				    DataaVecTextForMultipleChoice,
				    DataaVecTextForQuestionAnswering,
				    DataaVecTextForSequenceClassification,
				    DataaVecTextForTokenClassification,
				    DataaVecTextModel,
				    DataaVecTextPreTrainedModel,
				)
				from .modeling_dataavec_vision import (
				    DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DataaVecVisionForImageClassification,
				    DataaVecVisionForMaskedImageModeling,
				    DataaVecVisionForSemanticSegmentation,
				    DataaVecVisionModel,
				    DataaVecVisionPreTrainedModel,
				)
		if is_tf_available():
				from .modeling_tf_dataavec_vision import (
				    TFDataaVecVisionForImageClassification,
				    TFDataaVecVisionForSemanticSegmentation,
				    TFDataaVecVisionModel,
				    TFDataaVecVisionPreTrainedModel,
				)

else:
		import sys

		lowerCamelCase      				=     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






