

"""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							):







					'''simple docstring'''




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







					def  _lowerCamelCase		(						self	,			_UpperCAmelCase     ):
									return self.lineara(self.batchnorm(self.lineara(_UpperCAmelCase     )     )     )






class 				__lowercase			(			_UpperCamelCase							):







					'''simple docstring'''




					def  _lowerCamelCase		(						self	,			_UpperCAmelCase	,			*_UpperCAmelCase	,			**_UpperCAmelCase     ):
									return (args[0] + 1,) + args[1:], kwargs






class 				__lowercase			(			_UpperCamelCase							):







					'''simple docstring'''




					def  _lowerCamelCase		(						self	,			_UpperCAmelCase	,			_UpperCAmelCase     ):
									return output + 1






class 				__lowercase			(			unittest.TestCase							):







					'''simple docstring'''




					def  _lowerCamelCase		(						self     ):
									__a :						Dict			       =  ModelForTest()
									__a :						Dict			       =  ModelHook()

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







					def  _lowerCamelCase		(						self     ):
									__a :						List[str]			       =  ModelForTest()
									__a :						List[str]			       =  ModelHook()

									add_hook_to_module(_UpperCAmelCase	,			_UpperCAmelCase     )
									add_hook_to_module(_UpperCAmelCase	,			_UpperCAmelCase	,			append=_UpperCAmelCase     )

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







					def  _lowerCamelCase		(						self     ):
									__a :						Any			       =  ModelForTest()
									__a :						List[Any]			       =  torch.randn(2	,			3     )
									__a :						Optional[Any]			       =  test_model(x + 1     )
									__a :						Optional[int]			       =  test_model(x + 2     )

									__a :						Optional[int]			       =  PreForwardHook()
									add_hook_to_module(_UpperCAmelCase	,			_UpperCAmelCase     )
									__a :						Tuple			       =  test_model(_UpperCAmelCase     )
									self.assertTrue(torch.allclose(_UpperCAmelCase	,			_UpperCAmelCase	,			atol=1e-5     )     )

									# Attaching a hook to a model when it already has one replaces, does not chain
									__a :						Union[str, Any]			       =  PreForwardHook()
									add_hook_to_module(_UpperCAmelCase	,			_UpperCAmelCase     )
									__a :						Optional[int]			       =  test_model(_UpperCAmelCase     )
									self.assertTrue(torch.allclose(_UpperCAmelCase	,			_UpperCAmelCase	,			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(_UpperCAmelCase	,			_UpperCAmelCase     )

									__a :						List[str]			       =  test_model(_UpperCAmelCase     )
									assert torch.allclose(_UpperCAmelCase	,			_UpperCAmelCase	,			atol=1e-5     )







					def  _lowerCamelCase		(						self     ):
									__a :						str			       =  ModelForTest()
									__a :						Tuple			       =  torch.randn(2	,			3     )
									__a :						List[Any]			       =  test_model(_UpperCAmelCase     )

									__a :						int			       =  PostForwardHook()
									add_hook_to_module(_UpperCAmelCase	,			_UpperCAmelCase     )
									__a :						List[str]			       =  test_model(_UpperCAmelCase     )
									self.assertTrue(torch.allclose(_UpperCAmelCase	,			output + 1	,			atol=1e-5     )     )

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

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

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







					def  _lowerCamelCase		(						self     ):
									__a :						int			       =  ModelForTest()
									__a :						Optional[int]			       =  torch.randn(2	,			3     )
									__a :						int			       =  test_model(_UpperCAmelCase     )

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

									__a :						Union[str, Any]			       =  True
									__a :						List[Any]			       =  test_model(_UpperCAmelCase     )
									self.assertFalse(outputa.requires_grad     )







					@require_multi_gpu
					def  _lowerCamelCase		(						self     ):
									__a :						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 :						Dict			       =  torch.randn(2	,			3     )
									__a :						Dict			       =  model(_UpperCAmelCase     )
									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(_UpperCAmelCase	,			AlignDevicesHook(io_same_device=_UpperCAmelCase     )     )
									__a :						Union[str, Any]			       =  torch.randn(2	,			3     ).to(0     )
									__a :						Optional[int]			       =  model(_UpperCAmelCase     )
									self.assertEqual(output.device	,			torch.device(0     )     )







					def  _lowerCamelCase		(						self     ):
									__a :						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 :						int			       =  {'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

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

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

									__a :						Tuple			       =  torch.randn(2	,			3     )
									__a :						Optional[int]			       =  model(_UpperCAmelCase     )
									self.assertEqual(output.device	,			_UpperCAmelCase     )

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

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

									# 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 :						int			       =  torch.randn(2	,			3     )
									__a :						Tuple			       =  model(_UpperCAmelCase     )
									self.assertEqual(output.device	,			_UpperCAmelCase     )

									# 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     ):
									__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
									__a :						str			       =  0 if torch.cuda.is_available() else '''cpu'''
									attach_align_device_hook(_UpperCAmelCase	,			execution_device=_UpperCAmelCase	,			offload=_UpperCAmelCase     )

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

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

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

									# 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 :						str			       =  model(_UpperCAmelCase     )
									self.assertEqual(output.device	,			_UpperCAmelCase     )

									# Removing hooks loads back the weights in the model.
									remove_hook_from_submodules(_UpperCAmelCase     )
									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     ):
									__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 :						List[Any]			       =  0 if torch.cuda.is_available() else '''cpu'''
									attach_align_device_hook(
									    _UpperCAmelCase	,			execution_device=_UpperCAmelCase	,			offload=_UpperCAmelCase	,			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 :						List[Any]			       =  torch.device(_UpperCAmelCase     )
									self.assertEqual(model.batchnorm.running_mean.device	,			_UpperCAmelCase     )

									__a :						Dict			       =  torch.randn(2	,			3     )
									__a :						Union[str, Any]			       =  model(_UpperCAmelCase     )
									self.assertEqual(output.device	,			_UpperCAmelCase     )

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

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

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


A 			=  {
    '''configuration_jukebox''': [
        '''JUKEBOX_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''JukeboxConfig''',
        '''JukeboxPriorConfig''',
        '''JukeboxVQVAEConfig''',
    ],
    '''tokenization_jukebox''': ['''JukeboxTokenizer'''],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A 			=  [
		    '''JUKEBOX_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''JukeboxModel''',
		    '''JukeboxPreTrainedModel''',
		    '''JukeboxVQVAE''',
		    '''JukeboxPrior''',
		]

if TYPE_CHECKING:
		from .configuration_jukebox import (
		    JUKEBOX_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    JukeboxConfig,
		    JukeboxPriorConfig,
		    JukeboxVQVAEConfig,
		)
		from .tokenization_jukebox import JukeboxTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_jukebox import (
				    JUKEBOX_PRETRAINED_MODEL_ARCHIVE_LIST,
				    JukeboxModel,
				    JukeboxPreTrainedModel,
				    JukeboxPrior,
				    JukeboxVQVAE,
				)

else:
		import sys

		A 			=  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)