





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





      def __init__(			self					)						->	Optional[Any]:
            super().__init__()
            lowerCAmelCase    		=			nn.Linear(3 ,		4					)
            lowerCAmelCase    		=			nn.BatchNormad(4					)
            lowerCAmelCase    		=			nn.Linear(4 ,		5					)





      def        _snake_case     (			self ,		lowercase					)						->	Any:
            return self.lineara(self.batchnorm(self.lineara(lowercase					)					)					)







class 						lowercase   (   _UpperCAmelCase   ):





      def        _snake_case     (			self ,		lowercase ,		*lowercase ,		**lowercase					)						->	Optional[int]:
            return (args[0] + 1,) + args[1:], kwargs







class 						lowercase   (   _UpperCAmelCase   ):





      def        _snake_case     (			self ,		lowercase ,		lowercase					)						->	str:
            return output + 1







class 						lowercase   (   unittest.TestCase   ):





      def        _snake_case     (			self					)						->	Tuple:
            lowerCAmelCase    		=			ModelForTest()
            lowerCAmelCase    		=			ModelHook()

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





      def        _snake_case     (			self					)						->	Dict:
            lowerCAmelCase    		=			ModelForTest()
            lowerCAmelCase    		=			ModelHook()

            add_hook_to_module(lowercase ,		lowercase					)
            add_hook_to_module(lowercase ,		lowercase ,		append=lowercase					)

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





      def        _snake_case     (			self					)						->	str:
            lowerCAmelCase    		=			ModelForTest()
            lowerCAmelCase    		=			torch.randn(2 ,		3					)
            lowerCAmelCase    		=			test_model(x + 1					)
            lowerCAmelCase    		=			test_model(x + 2					)

            lowerCAmelCase    		=			PreForwardHook()
            add_hook_to_module(lowercase ,		lowercase					)
            lowerCAmelCase    		=			test_model(lowercase					)
            self.assertTrue(torch.allclose(lowercase ,		lowercase ,		atol=1e-5					)					)

            # Attaching a hook to a model when it already has one replaces, does not chain
            lowerCAmelCase    		=			PreForwardHook()
            add_hook_to_module(lowercase ,		lowercase					)
            lowerCAmelCase    		=			test_model(lowercase					)
            self.assertTrue(torch.allclose(lowercase ,		lowercase ,		atol=1e-5					)					)

            # You need to use the sequential hook to chain two or more hooks
            lowerCAmelCase    		=			SequentialHook(PreForwardHook() ,		PreForwardHook()					)
            add_hook_to_module(lowercase ,		lowercase					)

            lowerCAmelCase    		=			test_model(lowercase					)
            assert torch.allclose(lowercase ,		lowercase ,		atol=1e-5					)





      def        _snake_case     (			self					)						->	Optional[Any]:
            lowerCAmelCase    		=			ModelForTest()
            lowerCAmelCase    		=			torch.randn(2 ,		3					)
            lowerCAmelCase    		=			test_model(lowercase					)

            lowerCAmelCase    		=			PostForwardHook()
            add_hook_to_module(lowercase ,		lowercase					)
            lowerCAmelCase    		=			test_model(lowercase					)
            self.assertTrue(torch.allclose(lowercase ,		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(lowercase ,		lowercase					)
            lowerCAmelCase    		=			test_model(lowercase					)
            self.assertTrue(torch.allclose(lowercase ,		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(lowercase ,		lowercase					)

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





      def        _snake_case     (			self					)						->	Dict:
            lowerCAmelCase    		=			ModelForTest()
            lowerCAmelCase    		=			torch.randn(2 ,		3					)
            lowerCAmelCase    		=			test_model(lowercase					)

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

            lowerCAmelCase    		=			True
            lowerCAmelCase    		=			test_model(lowercase					)
            self.assertFalse(outputa.requires_grad					)





      @require_multi_gpu
      def        _snake_case     (			self					)						->	Dict:
            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(lowercase					)
            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(lowercase ,		AlignDevicesHook(io_same_device=lowercase					)					)
            lowerCAmelCase    		=			torch.randn(2 ,		3					).to(0					)
            lowerCAmelCase    		=			model(lowercase					)
            self.assertEqual(output.device ,		torch.device(0					)					)





      def        _snake_case     (			self					)						->	List[str]:
            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(**lowercase					)					)
            add_hook_to_module(model.batchnorm ,		AlignDevicesHook(**lowercase					)					)
            add_hook_to_module(model.lineara ,		AlignDevicesHook(**lowercase					)					)

            # 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 ,		lowercase					)

            lowerCAmelCase    		=			torch.randn(2 ,		3					)
            lowerCAmelCase    		=			model(lowercase					)
            self.assertEqual(output.device ,		lowercase					)

            # 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(**lowercase					)					)
            add_hook_to_module(model.batchnorm ,		AlignDevicesHook(**lowercase					)					)
            add_hook_to_module(model.lineara ,		AlignDevicesHook(**lowercase					)					)

            # 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(lowercase					)
            self.assertEqual(output.device ,		lowercase					)

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





      def        _snake_case     (			self					)						->	Union[str, Any]:
            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(lowercase ,		execution_device=lowercase ,		offload=lowercase					)

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

            lowerCAmelCase    		=			torch.randn(2 ,		3					)
            lowerCAmelCase    		=			model(lowercase					)
            self.assertEqual(output.device ,		lowercase					)

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

            # 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(lowercase					)
            self.assertEqual(output.device ,		lowercase					)

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





      def        _snake_case     (			self					)						->	str:
            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(
                lowercase ,		execution_device=lowercase ,		offload=lowercase ,		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(lowercase					)
            self.assertEqual(model.batchnorm.running_mean.device ,		lowercase					)

            lowerCAmelCase    		=			torch.randn(2 ,		3					)
            lowerCAmelCase    		=			model(lowercase					)
            self.assertEqual(output.device ,		lowercase					)

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

            # 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(lowercase					)
            self.assertEqual(output.device ,		lowercase					)

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


SCREAMING_SNAKE_CASE__   =  {
    "configuration_roformer": ["ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "RoFormerConfig", "RoFormerOnnxConfig"],
    "tokenization_roformer": ["RoFormerTokenizer"],
}

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  ["RoFormerTokenizerFast"]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "RoFormerForCausalLM",
           "RoFormerForMaskedLM",
           "RoFormerForMultipleChoice",
           "RoFormerForQuestionAnswering",
           "RoFormerForSequenceClassification",
           "RoFormerForTokenClassification",
           "RoFormerLayer",
           "RoFormerModel",
           "RoFormerPreTrainedModel",
           "load_tf_weights_in_roformer",
       ]


try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TFRoFormerForCausalLM",
           "TFRoFormerForMaskedLM",
           "TFRoFormerForMultipleChoice",
           "TFRoFormerForQuestionAnswering",
           "TFRoFormerForSequenceClassification",
           "TFRoFormerForTokenClassification",
           "TFRoFormerLayer",
           "TFRoFormerModel",
           "TFRoFormerPreTrainedModel",
       ]


try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "FlaxRoFormerForMaskedLM",
           "FlaxRoFormerForMultipleChoice",
           "FlaxRoFormerForQuestionAnswering",
           "FlaxRoFormerForSequenceClassification",
           "FlaxRoFormerForTokenClassification",
           "FlaxRoFormerModel",
           "FlaxRoFormerPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_roformer import ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, RoFormerConfig, RoFormerOnnxConfig
       from .tokenization_roformer import RoFormerTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_roformer_fast import RoFormerTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_roformer import (
                  ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  RoFormerForCausalLM,
                  RoFormerForMaskedLM,
                  RoFormerForMultipleChoice,
                  RoFormerForQuestionAnswering,
                  RoFormerForSequenceClassification,
                  RoFormerForTokenClassification,
                  RoFormerLayer,
                  RoFormerModel,
                  RoFormerPreTrainedModel,
                  load_tf_weights_in_roformer,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_roformer import (
                  TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFRoFormerForCausalLM,
                  TFRoFormerForMaskedLM,
                  TFRoFormerForMultipleChoice,
                  TFRoFormerForQuestionAnswering,
                  TFRoFormerForSequenceClassification,
                  TFRoFormerForTokenClassification,
                  TFRoFormerLayer,
                  TFRoFormerModel,
                  TFRoFormerPreTrainedModel,
              )

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_flax_roformer import (
                  FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  FlaxRoFormerForMaskedLM,
                  FlaxRoFormerForMultipleChoice,
                  FlaxRoFormerForQuestionAnswering,
                  FlaxRoFormerForSequenceClassification,
                  FlaxRoFormerForTokenClassification,
                  FlaxRoFormerModel,
                  FlaxRoFormerPreTrainedModel,
              )


else:
       import sys

       SCREAMING_SNAKE_CASE__   =  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






