





'''simple docstring'''






A   : Tuple             =			[
    'Audio',
    'Array2D',
    'Array3D',
    'Array4D',
    'Array5D',
    'ClassLabel',
    'Features',
    'Sequence',
    'Value',
    'Image',
    'Translation',
    'TranslationVariableLanguages',
]
from .audio import Audio
from .features import ArrayaD, ArrayaD, ArrayaD, ArrayaD, ClassLabel, Features, Sequence, Value
from .image import Image
from .translation import Translation, TranslationVariableLanguages












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


					def __init__(					self	:  Optional[int]       ):





											'''simple docstring'''





											super().__init__()
											_snake_case:   int  					=		nn.Linear(3       ,     4       )
											_snake_case:   Any  					=		nn.BatchNormad(4       )
											_snake_case:   Union[str, Any]  					=		nn.Linear(4       ,     5       )




					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  List[Any]       ):





											'''simple docstring'''





											return self.lineara(self.batchnorm(self.lineara(__snake_case       )       )       )







class 			lowerCamelCase					(						__UpperCAmelCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  Optional[Any]       ,     *__snake_case	:  Optional[Any]       ,     **__snake_case	:  List[Any]       ):





											'''simple docstring'''





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







class 			lowerCamelCase					(						__UpperCAmelCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Tuple       ):





											'''simple docstring'''





											return output + 1







class 			lowerCamelCase					(						unittest.TestCase  ):


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





											'''simple docstring'''





											_snake_case:   str  					=		ModelForTest()
											_snake_case:   List[str]  					=		ModelHook()

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


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		ModelForTest()
											_snake_case:   Union[str, Any]  					=		ModelHook()

											add_hook_to_module(__snake_case       ,     __snake_case       )
											add_hook_to_module(__snake_case       ,     __snake_case       ,     append=__snake_case       )

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


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





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

											_snake_case:   List[str]  					=		PreForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   List[Any]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )       )

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

											# You need to use the sequential hook to chain two or more hooks
											_snake_case:   Any  					=		SequentialHook(PreForwardHook()       ,     PreForwardHook()       )
											add_hook_to_module(__snake_case       ,     __snake_case       )

											_snake_case:   int  					=		test_model(__snake_case       )
											assert torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   int  					=		ModelForTest()
											_snake_case:   str  					=		torch.randn(2       ,     3       )
											_snake_case:   int  					=		test_model(__snake_case       )

											_snake_case:   Union[str, Any]  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   int  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       ,     atol=1e-5       )       )

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

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

											_snake_case:   List[Any]  					=		test_model(__snake_case       )
											assert torch.allclose(__snake_case       ,     output + 2       ,     atol=1e-5       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		ModelForTest()
											_snake_case:   Dict  					=		torch.randn(2       ,     3       )
											_snake_case:   Tuple  					=		test_model(__snake_case       )

											_snake_case:   Tuple  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   List[str]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       )       )
											self.assertTrue(outputa.requires_grad       )

											_snake_case:   Any  					=		True
											_snake_case:   str  					=		test_model(__snake_case       )
											self.assertFalse(outputa.requires_grad       )


					@require_multi_gpu
					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   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
											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:   int  					=		torch.randn(2       ,     3       )
											_snake_case:   Tuple  					=		model(__snake_case       )
											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(__snake_case       ,     AlignDevicesHook(io_same_device=__snake_case       )       )
											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       ).to(0       )
											_snake_case:   Any  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     torch.device(0       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





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

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

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

											_snake_case:   Tuple  					=		torch.randn(2       ,     3       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

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

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

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

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





											'''simple docstring'''





											_snake_case:   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
											_snake_case:   Optional[Any]  					=		0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(__snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       )

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

											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       )
											_snake_case:   List[str]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

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

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

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





											'''simple docstring'''





											_snake_case:   Tuple  					=		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:   Tuple  					=		0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(
											    __snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       ,     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:   str  					=		torch.device(__snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

											_snake_case:   Dict  					=		torch.randn(2       ,     3       )
											_snake_case:   Dict  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

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

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

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





