





'''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    ) ->	str:
      super().__init__()
      snake_case				=			nn.Linear(3    ,   4    )
      snake_case				=			nn.BatchNormad(4    )
      snake_case				=			nn.Linear(4    ,   5    )







     def     UpperCamelCase			( self    ,   A__    ) ->	Tuple:
      return self.lineara(self.batchnorm(self.lineara(lowerCamelCase__    )    )    )




class 			_lowercase						(       __lowerCAmelCase					):







     def     UpperCamelCase			( self    ,   A__    ,   *A__    ,   **A__    ) ->	Any:
      return (args[0] + 1,) + args[1:], kwargs




class 			_lowercase						(       __lowerCAmelCase					):







     def     UpperCamelCase			( self    ,   A__    ,   A__    ) ->	Any:
      return output + 1




class 			_lowercase						(       unittest.TestCase					):







     def     UpperCamelCase			( self    ) ->	Tuple:
      snake_case				=			ModelForTest()
      snake_case				=			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    ) ->	str:
      snake_case				=			ModelForTest()
      snake_case				=			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    ) ->	int:
      snake_case				=			ModelForTest()
      snake_case				=			torch.randn(2    ,   3    )
      snake_case				=			test_model(x + 1    )
      snake_case				=			test_model(x + 2    )

      snake_case				=			PreForwardHook()
      add_hook_to_module(lowerCamelCase__    ,   lowerCamelCase__    )
      snake_case				=			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
      snake_case				=			PreForwardHook()
      add_hook_to_module(lowerCamelCase__    ,   lowerCamelCase__    )
      snake_case				=			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
      snake_case				=			SequentialHook(PreForwardHook()    ,   PreForwardHook()    )
      add_hook_to_module(lowerCamelCase__    ,   lowerCamelCase__    )

      snake_case				=			test_model(lowerCamelCase__    )
      assert torch.allclose(lowerCamelCase__    ,   lowerCamelCase__    ,   atol=1e-5    )







     def     UpperCamelCase			( self    ) ->	Optional[int]:
      snake_case				=			ModelForTest()
      snake_case				=			torch.randn(2    ,   3    )
      snake_case				=			test_model(lowerCamelCase__    )

      snake_case				=			PostForwardHook()
      add_hook_to_module(lowerCamelCase__    ,   lowerCamelCase__    )
      snake_case				=			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
      snake_case				=			PostForwardHook()
      add_hook_to_module(lowerCamelCase__    ,   lowerCamelCase__    )
      snake_case				=			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
      snake_case				=			SequentialHook(PostForwardHook()    ,   PostForwardHook()    )
      add_hook_to_module(lowerCamelCase__    ,   lowerCamelCase__    )

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







     def     UpperCamelCase			( self    ) ->	Dict:
      snake_case				=			ModelForTest()
      snake_case				=			torch.randn(2    ,   3    )
      snake_case				=			test_model(lowerCamelCase__    )

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

      snake_case				=			True
      snake_case				=			test_model(lowerCamelCase__    )
      self.assertFalse(outputa.requires_grad    )







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







     def     UpperCamelCase			( self    ) ->	Union[str, Any]:
      snake_case				=			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
      snake_case				=			{'''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
      snake_case				=			torch.device(hook_kwargs['''execution_device''']    )
      self.assertEqual(model.batchnorm.running_mean.device    ,   lowerCamelCase__    )

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

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

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

      snake_case				=			torch.randn(2    ,   3    )
      snake_case				=			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    ) ->	List[Any]:
      snake_case				=			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
      snake_case				=			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
      snake_case				=			torch.device(lowerCamelCase__    )
      self.assertEqual(model.batchnorm.running_mean.device    ,   lowerCamelCase__    )

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

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








'''simple docstring'''







from ..utils import DummyObject, requires_backends
class 			_lowercase						(       metaclass=__a					):
     _UpperCAmelCase										=					['''transformers''', '''torch''', '''note_seq''']







     def __init__( self    ,   *A__    ,   **A__    ) ->	Union[str, Any]:
      requires_backends(self    ,   ['''transformers''', '''torch''', '''note_seq''']    )







     @classmethod
     def     UpperCamelCase			( cls    ,   *A__    ,   **A__    ) ->	Optional[Any]:
      requires_backends(cls    ,   ['''transformers''', '''torch''', '''note_seq''']    )







     @classmethod
     def     UpperCamelCase			( cls    ,   *A__    ,   **A__    ) ->	Any:
      requires_backends(cls    ,   ['''transformers''', '''torch''', '''note_seq''']    )

