






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		:    Union[str, Any]    ):



					'''simple docstring'''



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



		def    __lowerCAmelCase  (					self		:    List[str]    ,       SCREAMING_SNAKE_CASE__		:    Union[str, Any]    ):



					'''simple docstring'''



					return self.lineara(self.batchnorm(self.lineara(__lowerCamelCase    )    )    )






class 			_UpperCamelCase(       __SCREAMING_SNAKE_CASE		):

		def    __lowerCAmelCase  (					self		:    Optional[Any]    ,       SCREAMING_SNAKE_CASE__		:    Optional[int]    ,       *SCREAMING_SNAKE_CASE__		:    List[Any]    ,       **SCREAMING_SNAKE_CASE__		:    Union[str, Any]    ):



					'''simple docstring'''



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






class 			_UpperCamelCase(       __SCREAMING_SNAKE_CASE		):

		def    __lowerCAmelCase  (					self		:    int    ,       SCREAMING_SNAKE_CASE__		:    Union[str, Any]    ,       SCREAMING_SNAKE_CASE__		:    str    ):



					'''simple docstring'''



					return output + 1






class 			_UpperCamelCase(       unittest.TestCase		):

		def    __lowerCAmelCase  (					self		:    Tuple    ):



					'''simple docstring'''



					__a						:  Optional[Any]									=     ModelForTest()
					__a						:  List[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    __lowerCAmelCase  (					self		:    Tuple    ):



					'''simple docstring'''



					__a						:  str									=     ModelForTest()
					__a						:  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    __lowerCAmelCase  (					self		:    Any    ):



					'''simple docstring'''



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

					__a						:  Optional[Any]									=     PreForwardHook()
					add_hook_to_module(__lowerCamelCase    ,       __lowerCamelCase    )
					__a						:  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
					__a						:  Tuple									=     PreForwardHook()
					add_hook_to_module(__lowerCamelCase    ,       __lowerCamelCase    )
					__a						:  Any									=     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
					__a						:  List[str]									=     SequentialHook(PreForwardHook()    ,       PreForwardHook()    )
					add_hook_to_module(__lowerCamelCase    ,       __lowerCamelCase    )

					__a						:  Union[str, Any]									=     test_model(__lowerCamelCase    )
					assert torch.allclose(__lowerCamelCase    ,       __lowerCamelCase    ,       atol=1e-5    )

		def    __lowerCAmelCase  (					self		:    Union[str, Any]    ):



					'''simple docstring'''



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

					__a						:  Union[str, Any]									=     PostForwardHook()
					add_hook_to_module(__lowerCamelCase    ,       __lowerCamelCase    )
					__a						:  Union[str, Any]									=     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
					__a						:  Tuple									=     PostForwardHook()
					add_hook_to_module(__lowerCamelCase    ,       __lowerCamelCase    )
					__a						:  Dict									=     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
					__a						:  str									=     SequentialHook(PostForwardHook()    ,       PostForwardHook()    )
					add_hook_to_module(__lowerCamelCase    ,       __lowerCamelCase    )

					__a						:  Dict									=     test_model(__lowerCamelCase    )
					assert torch.allclose(__lowerCamelCase    ,       output + 2    ,       atol=1e-5    )

		def    __lowerCAmelCase  (					self		:    Any    ):



					'''simple docstring'''



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

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

					__a						:  Tuple									=     True
					__a						:  Optional[Any]									=     test_model(__lowerCamelCase    )
					self.assertFalse(outputa.requires_grad    )

		@require_multi_gpu
		def    __lowerCAmelCase  (					self		:    List[str]    ):



					'''simple docstring'''



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

		def    __lowerCAmelCase  (					self		:    Tuple    ):



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

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

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



					'''simple docstring'''



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

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

					# 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'    )    )

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



		def    __lowerCAmelCase  (					self		:    Dict    ):



					'''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						:  Tuple									=     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
					__a						:  Union[str, Any]									=     torch.device(__lowerCamelCase    )
					self.assertEqual(model.batchnorm.running_mean.device    ,       __lowerCamelCase    )

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

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










"""simple docstring"""






def 			snake_case		(		lowerCAmelCase_  ,						lowerCAmelCase_	)						->				int:
  return 1 if input_a == input_a else 0




def 			snake_case		(		)						->				None:
  assert xnor_gate(0  ,						0	) == 1
  assert xnor_gate(0  ,						1	) == 0
  assert xnor_gate(1  ,						0	) == 0
  assert xnor_gate(1  ,						1	) == 1


if __name__ == "__main__":
     print(xnor_gate(0, 0))
     print(xnor_gate(0, 1))
     print(xnor_gate(1, 0))
     print(xnor_gate(1, 1))





