


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


					def __init__(      self							:     Dict			)  ->      str:





						'''simple docstring'''






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


					def      UpperCAmelCase__	(      self							:     int							,	lowerCamelCase__							:     Dict			)  ->      Tuple:





						'''simple docstring'''






						return self.lineara(self.batchnorm(self.lineara(lowerCamelCase__			)			)			)



class       __magic_name__	(     __lowerCAmelCase):


					def      UpperCAmelCase__	(      self							:     Union[str, Any]							,	lowerCamelCase__							:     Any							,	*lowerCamelCase__							:     int							,	**lowerCamelCase__							:     List[str]			)  ->      Any:





						'''simple docstring'''






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



class       __magic_name__	(     __lowerCAmelCase):


					def      UpperCAmelCase__	(      self							:     Any							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     Any			)  ->      Any:





						'''simple docstring'''






						return output + 1



class       __magic_name__	(     unittest.TestCase):


					def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Tuple:





						'''simple docstring'''






						UpperCamelCase__					:    Tuple          =   ModelForTest()
						UpperCamelCase__					:    Union[str, Any]          =   ModelHook()

						add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
						self.assertEqual(test_model._hf_hook							,	lowerCamelCase__			)
						self.assertTrue(hasattr(lowerCamelCase__							,	'''_old_forward'''			)			)

						# Check adding the hook did not change the name or the signature
						self.assertEqual(test_model.forward.__name__							,	'''forward'''			)
						self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)							,	['''x''']			)

						remove_hook_from_module(lowerCamelCase__			)
						self.assertFalse(hasattr(lowerCamelCase__							,	'''_hf_hook'''			)			)
						self.assertFalse(hasattr(lowerCamelCase__							,	'''_old_forward'''			)			)


					def      UpperCAmelCase__	(      self							:     Dict			)  ->      str:





						'''simple docstring'''






						UpperCamelCase__					:    Dict          =   ModelForTest()
						UpperCamelCase__					:    List[str]          =   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      UpperCAmelCase__	(      self							:     List[str]			)  ->      int:





						'''simple docstring'''






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

						UpperCamelCase__					:    Optional[int]          =   PreForwardHook()
						add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
						UpperCamelCase__					:    List[Any]          =   test_model(lowerCamelCase__			)
						self.assertTrue(torch.allclose(lowerCamelCase__							,	lowerCamelCase__							,	atol=1E-5			)			)

						# Attaching a hook to a model when it already has one replaces, does not chain
						UpperCamelCase__					:    Optional[Any]          =   PreForwardHook()
						add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
						UpperCamelCase__					:    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
						UpperCamelCase__					:    Tuple          =   SequentialHook(PreForwardHook()							,	PreForwardHook()			)
						add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)

						UpperCamelCase__					:    Any          =   test_model(lowerCamelCase__			)
						assert torch.allclose(lowerCamelCase__							,	lowerCamelCase__							,	atol=1E-5			)


					def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      Optional[int]:





						'''simple docstring'''






						UpperCamelCase__					:    Dict          =   ModelForTest()
						UpperCamelCase__					:    str          =   torch.randn(2							,	3			)
						UpperCamelCase__					:    Any          =   test_model(lowerCamelCase__			)

						UpperCamelCase__					:    int          =   PostForwardHook()
						add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
						UpperCamelCase__					:    Dict          =   test_model(lowerCamelCase__			)
						self.assertTrue(torch.allclose(lowerCamelCase__							,	output + 1							,	atol=1E-5			)			)

						# Attaching a hook to a model when it already has one replaces, does not chain
						UpperCamelCase__					:    Dict          =   PostForwardHook()
						add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
						UpperCamelCase__					:    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
						UpperCamelCase__					:    List[Any]          =   SequentialHook(PostForwardHook()							,	PostForwardHook()			)
						add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)

						UpperCamelCase__					:    Tuple          =   test_model(lowerCamelCase__			)
						assert torch.allclose(lowerCamelCase__							,	output + 2							,	atol=1E-5			)


					def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      Dict:





						'''simple docstring'''






						UpperCamelCase__					:    List[Any]          =   ModelForTest()
						UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			)
						UpperCamelCase__					:    Tuple          =   test_model(lowerCamelCase__			)

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

						UpperCamelCase__					:    Optional[int]          =   True
						UpperCamelCase__					:    Optional[int]          =   test_model(lowerCamelCase__			)
						self.assertFalse(outputa.requires_grad			)


					@require_multi_gpu
					def      UpperCAmelCase__	(      self							:     Any			)  ->      Dict:





						'''simple docstring'''






						UpperCamelCase__					:    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
						UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			)
						UpperCamelCase__					:    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__			)			)
						UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			).to(0			)
						UpperCamelCase__					:    Optional[int]          =   model(lowerCamelCase__			)
						self.assertEqual(output.device							,	torch.device(0			)			)


					def      UpperCAmelCase__	(      self							:     List[str]			)  ->      Union[str, Any]:





						'''simple docstring'''






						UpperCamelCase__					:    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
						UpperCamelCase__					:    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
						UpperCamelCase__					:    int          =   torch.device(hook_kwargs['''execution_device''']			)
						self.assertEqual(model.batchnorm.running_mean.device							,	lowerCamelCase__			)

						UpperCamelCase__					:    Dict          =   torch.randn(2							,	3			)
						UpperCamelCase__					:    List[str]          =   model(lowerCamelCase__			)
						self.assertEqual(output.device							,	lowerCamelCase__			)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_module(model.lineara			)
						remove_hook_from_module(model.batchnorm			)
						remove_hook_from_module(model.lineara			)
						self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
						self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
						self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

						# Now test with buffers included in the offload
						UpperCamelCase__					:    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'''			)			)

						UpperCamelCase__					:    List[Any]          =   torch.randn(2							,	3			)
						UpperCamelCase__					:    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      UpperCAmelCase__	(      self							:     str			)  ->      Optional[int]:





						'''simple docstring'''






						UpperCamelCase__					:    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
						UpperCamelCase__					:    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
						UpperCamelCase__					:    str          =   torch.device(lowerCamelCase__			)
						self.assertEqual(model.batchnorm.running_mean.device							,	lowerCamelCase__			)

						UpperCamelCase__					:    int          =   torch.randn(2							,	3			)
						UpperCamelCase__					:    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__							,	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'''			)			)

						UpperCamelCase__					:    int          =   torch.randn(2							,	3			)
						UpperCamelCase__					:    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      UpperCAmelCase__	(      self							:     int			)  ->      List[Any]:





						'''simple docstring'''






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

						UpperCamelCase__					:    List[str]          =   torch.randn(2							,	3			)
						UpperCamelCase__					:    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'''			)			)

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






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


__UpperCamelCase       :       Optional[int]						       =				{
    "configuration_rag": ["RagConfig"],
    "retrieval_rag": ["RagRetriever"],
    "tokenization_rag": ["RagTokenizer"],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase       :       Optional[int]						       =				[
           "RagModel",
           "RagPreTrainedModel",
           "RagSequenceForGeneration",
           "RagTokenForGeneration",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase       :       List[Any]						       =				[
           "TFRagModel",
           "TFRagPreTrainedModel",
           "TFRagSequenceForGeneration",
           "TFRagTokenForGeneration",
       ]


if TYPE_CHECKING:
       from .configuration_rag import RagConfig
       from .retrieval_rag import RagRetriever
       from .tokenization_rag import RagTokenizer

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_rag import (
                  TFRagModel,
                  TFRagPreTrainedModel,
                  TFRagSequenceForGeneration,
                  TFRagTokenForGeneration,
              )

else:
       import sys

       __UpperCamelCase       :       List[Any]						       =				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


