


'''simple docstring'''



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




	def __init__(       self							):

						'''simple docstring'''





						super().__init__()
						__a      :	List[str]    =  nn.Linear(3						,    4							)
						__a      :	Optional[int]    =  nn.BatchNormad(4							)
						__a      :	int    =  nn.Linear(4						,    5							)





	def 				__UpperCAmelCase							(       self						,    __a							):

						'''simple docstring'''





						return self.lineara(self.batchnorm(self.lineara(__a							)							)							)





class 							__UpperCamelCase    (		lowerCAmelCase_ ):




	def 				__UpperCAmelCase							(       self						,    __a						,    *__a						,    **__a							):

						'''simple docstring'''





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





class 							__UpperCamelCase    (		lowerCAmelCase_ ):




	def 				__UpperCAmelCase							(       self						,    __a						,    __a							):

						'''simple docstring'''





						return output + 1





class 							__UpperCamelCase    (		unittest.TestCase ):




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Optional[int]    =  ModelForTest()
						__a      :	Any    =  ModelHook()

						add_hook_to_module(__a						,    __a							)
						self.assertEqual(test_model._hf_hook						,    __a							)
						self.assertTrue(hasattr(__a						,    '_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(__a							)
						self.assertFalse(hasattr(__a						,    '_hf_hook'							)							)
						self.assertFalse(hasattr(__a						,    '_old_forward'							)							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	int    =  ModelForTest()
						__a      :	int    =  ModelHook()

						add_hook_to_module(__a						,    __a							)
						add_hook_to_module(__a						,    __a						,    append=__a							)

						self.assertEqual(isinstance(test_model._hf_hook						,    __a							)						,    __a							)
						self.assertEqual(len(test_model._hf_hook.hooks							)						,    2							)
						self.assertTrue(hasattr(__a						,    '_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(__a							)
						self.assertFalse(hasattr(__a						,    '_hf_hook'							)							)
						self.assertFalse(hasattr(__a						,    '_old_forward'							)							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Any    =  ModelForTest()
						__a      :	Optional[Any]    =  torch.randn(2						,    3							)
						__a      :	str    =  test_model(x + 1							)
						__a      :	List[Any]    =  test_model(x + 2							)

						__a      :	Optional[int]    =  PreForwardHook()
						add_hook_to_module(__a						,    __a							)
						__a      :	Tuple    =  test_model(__a							)
						self.assertTrue(torch.allclose(__a						,    __a						,    atol=1E-5							)							)

						# Attaching a hook to a model when it already has one replaces, does not chain
						__a      :	Optional[int]    =  PreForwardHook()
						add_hook_to_module(__a						,    __a							)
						__a      :	Optional[Any]    =  test_model(__a							)
						self.assertTrue(torch.allclose(__a						,    __a						,    atol=1E-5							)							)

						# You need to use the sequential hook to chain two or more hooks
						__a      :	Optional[Any]    =  SequentialHook(PreForwardHook()						,    PreForwardHook()							)
						add_hook_to_module(__a						,    __a							)

						__a      :	List[str]    =  test_model(__a							)
						assert torch.allclose(__a						,    __a						,    atol=1E-5							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Union[str, Any]    =  ModelForTest()
						__a      :	int    =  torch.randn(2						,    3							)
						__a      :	List[str]    =  test_model(__a							)

						__a      :	Dict    =  PostForwardHook()
						add_hook_to_module(__a						,    __a							)
						__a      :	Union[str, Any]    =  test_model(__a							)
						self.assertTrue(torch.allclose(__a						,    output + 1						,    atol=1E-5							)							)

						# Attaching a hook to a model when it already has one replaces, does not chain
						__a      :	Any    =  PostForwardHook()
						add_hook_to_module(__a						,    __a							)
						__a      :	Optional[Any]    =  test_model(__a							)
						self.assertTrue(torch.allclose(__a						,    output + 1						,    atol=1E-5							)							)

						# You need to use the sequential hook to chain two or more hooks
						__a      :	Union[str, Any]    =  SequentialHook(PostForwardHook()						,    PostForwardHook()							)
						add_hook_to_module(__a						,    __a							)

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




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	str    =  ModelForTest()
						__a      :	List[Any]    =  torch.randn(2						,    3							)
						__a      :	Union[str, Any]    =  test_model(__a							)

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

						__a      :	Dict    =  True
						__a      :	Tuple    =  test_model(__a							)
						self.assertFalse(outputa.requires_grad							)




	@require_multi_gpu
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__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      :	Union[str, Any]    =  torch.randn(2						,    3							)
						__a      :	Any    =  model(__a							)
						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(__a						,    AlignDevicesHook(io_same_device=__a							)							)
						__a      :	List[str]    =  torch.randn(2						,    3							).to(0							)
						__a      :	Any    =  model(__a							)
						self.assertEqual(output.device						,    torch.device(0							)							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	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
						__a      :	List[Any]    =  {'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

						add_hook_to_module(model.lineara						,    AlignDevicesHook(**__a							)							)
						add_hook_to_module(model.batchnorm						,    AlignDevicesHook(**__a							)							)
						add_hook_to_module(model.lineara						,    AlignDevicesHook(**__a							)							)

						# 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      :	List[Any]    =  torch.device(hook_kwargs['execution_device']							)
						self.assertEqual(model.batchnorm.running_mean.device						,    __a							)

						__a      :	Optional[int]    =  torch.randn(2						,    3							)
						__a      :	Tuple    =  model(__a							)
						self.assertEqual(output.device						,    __a							)

						# 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      :	List[Any]    =  {
						    'execution_device': 0 if torch.cuda.is_available() else 'cpu',
						    'offload': True,
						    'offload_buffers': True,
						}

						add_hook_to_module(model.lineara						,    AlignDevicesHook(**__a							)							)
						add_hook_to_module(model.batchnorm						,    AlignDevicesHook(**__a							)							)
						add_hook_to_module(model.lineara						,    AlignDevicesHook(**__a							)							)

						# 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      :	Tuple    =  torch.randn(2						,    3							)
						__a      :	Any    =  model(__a							)
						self.assertEqual(output.device						,    __a							)

						# 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 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__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      :	str    =  0 if torch.cuda.is_available() else 'cpu'
						attach_align_device_hook(__a						,    execution_device=__a						,    offload=__a							)

						# 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(__a							)
						self.assertEqual(model.batchnorm.running_mean.device						,    __a							)

						__a      :	int    =  torch.randn(2						,    3							)
						__a      :	List[Any]    =  model(__a							)
						self.assertEqual(output.device						,    __a							)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_submodules(__a							)
						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(__a						,    execution_device=__a						,    offload=__a						,    offload_buffers=__a							)

						# 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      :	Optional[int]    =  model(__a							)
						self.assertEqual(output.device						,    __a							)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_submodules(__a							)
						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 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Optional[int]    =  ModelForTest()

						# Everything is on CPU
						self.assertEqual(model.lineara.weight.device						,    torch.device('cpu'							)							)
						self.assertEqual(model.batchnorm.weight.device						,    torch.device('cpu'							)							)
						self.assertEqual(model.lineara.weight.device						,    torch.device('cpu'							)							)

						# This will move each submodule on different devices
						__a      :	Optional[int]    =  0 if torch.cuda.is_available() else 'cpu'
						attach_align_device_hook(
						    __a						,    execution_device=__a						,    offload=__a						,    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      :	Union[str, Any]    =  torch.device(__a							)
						self.assertEqual(model.batchnorm.running_mean.device						,    __a							)

						__a      :	Tuple    =  torch.randn(2						,    3							)
						__a      :	List[Any]    =  model(__a							)
						self.assertEqual(output.device						,    __a							)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_submodules(__a							)
						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(
						    __a						,    execution_device=__a						,    offload=__a						,    weights_map=model.state_dict()						,    offload_buffers=__a						,    )

						# 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      :	Optional[Any]    =  torch.randn(2						,    3							)
						__a      :	List[Any]    =  model(__a							)
						self.assertEqual(output.device						,    __a							)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_submodules(__a							)
						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'''



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowercase     :						Tuple	=     {
    'configuration_distilbert': [
        'DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'DistilBertConfig',
        'DistilBertOnnxConfig',
    ],
    'tokenization_distilbert': ['DistilBertTokenizer'],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowercase     :						str	=     ['DistilBertTokenizerFast']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowercase     :						Any	=     [
	    'DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'DistilBertForMaskedLM',
	    'DistilBertForMultipleChoice',
	    'DistilBertForQuestionAnswering',
	    'DistilBertForSequenceClassification',
	    'DistilBertForTokenClassification',
	    'DistilBertModel',
	    'DistilBertPreTrainedModel',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowercase     :						List[str]	=     [
	    'TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'TFDistilBertForMaskedLM',
	    'TFDistilBertForMultipleChoice',
	    'TFDistilBertForQuestionAnswering',
	    'TFDistilBertForSequenceClassification',
	    'TFDistilBertForTokenClassification',
	    'TFDistilBertMainLayer',
	    'TFDistilBertModel',
	    'TFDistilBertPreTrainedModel',
	]

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__lowercase     :						List[str]	=     [
	    'FlaxDistilBertForMaskedLM',
	    'FlaxDistilBertForMultipleChoice',
	    'FlaxDistilBertForQuestionAnswering',
	    'FlaxDistilBertForSequenceClassification',
	    'FlaxDistilBertForTokenClassification',
	    'FlaxDistilBertModel',
	    'FlaxDistilBertPreTrainedModel',
	]


if TYPE_CHECKING:
	from .configuration_distilbert import (
	    DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    DistilBertConfig,
	    DistilBertOnnxConfig,
	)
	from .tokenization_distilbert import DistilBertTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_distilbert_fast import DistilBertTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_distilbert import (
		    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DistilBertForMaskedLM,
		    DistilBertForMultipleChoice,
		    DistilBertForQuestionAnswering,
		    DistilBertForSequenceClassification,
		    DistilBertForTokenClassification,
		    DistilBertModel,
		    DistilBertPreTrainedModel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_distilbert import (
		    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFDistilBertForMaskedLM,
		    TFDistilBertForMultipleChoice,
		    TFDistilBertForQuestionAnswering,
		    TFDistilBertForSequenceClassification,
		    TFDistilBertForTokenClassification,
		    TFDistilBertMainLayer,
		    TFDistilBertModel,
		    TFDistilBertPreTrainedModel,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_flax_distilbert import (
		    FlaxDistilBertForMaskedLM,
		    FlaxDistilBertForMultipleChoice,
		    FlaxDistilBertForQuestionAnswering,
		    FlaxDistilBertForSequenceClassification,
		    FlaxDistilBertForTokenClassification,
		    FlaxDistilBertModel,
		    FlaxDistilBertPreTrainedModel,
		)

else:
	import sys

	__lowercase     :						int	=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



