



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      :  List[Any]       )		->     Tuple:
											super().__init__()
											A      						=    nn.Linear(3						, 4       )
											A      						=    nn.BatchNormad(4       )
											A      						=    nn.Linear(4						, 5       )


				def __UpperCamelCase						(							self      :  Union[str, Any]						, __UpperCamelCase      :  List[Any]       )		->     str:
											return self.lineara(self.batchnorm(self.lineara(__UpperCamelCase       )       )       )





class        lowerCAmelCase__					(  _lowerCamelCase       ):


				def __UpperCamelCase						(							self      :  Any						, __UpperCamelCase      :  Any						, *__UpperCamelCase      :  Any						, **__UpperCamelCase      :  List[str]       )		->     Dict:
											return (args[0] + 1,) + args[1:], kwargs





class        lowerCAmelCase__					(  _lowerCamelCase       ):


				def __UpperCamelCase						(							self      :  Tuple						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  Tuple       )		->     str:
											return output + 1





class        lowerCAmelCase__					(  unittest.TestCase       ):


				def __UpperCamelCase						(							self      :  Dict       )		->     Any:
											A      						=    ModelForTest()
											A      						=    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 __UpperCamelCase						(							self      :  Any       )		->     Union[str, Any]:
											A      						=    ModelForTest()
											A      						=    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 __UpperCamelCase						(							self      :  List[Any]       )		->     Optional[Any]:
											A      						=    ModelForTest()
											A      						=    torch.randn(2						, 3       )
											A      						=    test_model(x + 1       )
											A      						=    test_model(x + 2       )

											A      						=    PreForwardHook()
											add_hook_to_module(__UpperCamelCase						, __UpperCamelCase       )
											A      						=    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      						=    PreForwardHook()
											add_hook_to_module(__UpperCamelCase						, __UpperCamelCase       )
											A      						=    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      						=    SequentialHook(PreForwardHook()						, PreForwardHook()       )
											add_hook_to_module(__UpperCamelCase						, __UpperCamelCase       )

											A      						=    test_model(__UpperCamelCase       )
											assert torch.allclose(__UpperCamelCase						, __UpperCamelCase						, atol=1e-5       )


				def __UpperCamelCase						(							self      :  Optional[int]       )		->     Optional[Any]:
											A      						=    ModelForTest()
											A      						=    torch.randn(2						, 3       )
											A      						=    test_model(__UpperCamelCase       )

											A      						=    PostForwardHook()
											add_hook_to_module(__UpperCamelCase						, __UpperCamelCase       )
											A      						=    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      						=    PostForwardHook()
											add_hook_to_module(__UpperCamelCase						, __UpperCamelCase       )
											A      						=    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      						=    SequentialHook(PostForwardHook()						, PostForwardHook()       )
											add_hook_to_module(__UpperCamelCase						, __UpperCamelCase       )

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


				def __UpperCamelCase						(							self      :  Optional[int]       )		->     List[str]:
											A      						=    ModelForTest()
											A      						=    torch.randn(2						, 3       )
											A      						=    test_model(__UpperCamelCase       )

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

											A      						=    True
											A      						=    test_model(__UpperCamelCase       )
											self.assertFalse(outputa.requires_grad       )


				@require_multi_gpu
				def __UpperCamelCase						(							self      :  List[Any]       )		->     List[Any]:
											A      						=    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      						=    torch.randn(2						, 3       )
											A      						=    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      						=    torch.randn(2						, 3       ).to(0       )
											A      						=    model(__UpperCamelCase       )
											self.assertEqual(output.device						, torch.device(0       )       )


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

											A      						=    torch.randn(2						, 3       )
											A      						=    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      						=    {
											    '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      						=    torch.randn(2						, 3       )
											A      						=    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 __UpperCamelCase						(							self      :  Any       )		->     Optional[Any]:
											A      						=    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      						=    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      						=    torch.device(__UpperCamelCase       )
											self.assertEqual(model.batchnorm.running_mean.device						, __UpperCamelCase       )

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

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




from ...configuration_utils import PretrainedConfig
from ...utils import logging


__snake_case     :Optional[Any]			       =logging.get_logger(__name__)

__snake_case     :Tuple			       ={
    'transfo-xl-wt103': 'https://huggingface.co/transfo-xl-wt103/resolve/main/config.json',
}



class        lowerCAmelCase__					(  _lowerCamelCase       ):
				A_   :     Any								=     'transfo-xl'
				A_   :     Dict								=     ['mems']
				A_   :     int								=     {
				    'n_token': 'vocab_size',
				    'hidden_size': 'd_model',
				    'num_attention_heads': 'n_head',
				    'num_hidden_layers': 'n_layer',
				}


				def __init__(							self      :  Optional[Any]						, __UpperCamelCase      :  str=267_735						, __UpperCamelCase      :  int=[20_000, 40_000, 200_000]						, __UpperCamelCase      :  List[str]=1_024						, __UpperCamelCase      :  Union[str, Any]=1_024						, __UpperCamelCase      :  Optional[Any]=16						, __UpperCamelCase      :  Dict=64						, __UpperCamelCase      :  Optional[Any]=4_096						, __UpperCamelCase      :  int=4						, __UpperCamelCase      :  List[str]=False						, __UpperCamelCase      :  Optional[int]=18						, __UpperCamelCase      :  int=1_600						, __UpperCamelCase      :  Dict=1_000						, __UpperCamelCase      :  Optional[Any]=True						, __UpperCamelCase      :  Optional[Any]=True						, __UpperCamelCase      :  Optional[Any]=0						, __UpperCamelCase      :  Any=-1						, __UpperCamelCase      :  Any=True						, __UpperCamelCase      :  List[Any]=0.1						, __UpperCamelCase      :  str=0.0						, __UpperCamelCase      :  Dict=True						, __UpperCamelCase      :  int="normal"						, __UpperCamelCase      :  List[Any]=0.0_1						, __UpperCamelCase      :  str=0.0_1						, __UpperCamelCase      :  Any=0.0_2						, __UpperCamelCase      :  str=1e-5						, __UpperCamelCase      :  Optional[int]=0						, **__UpperCamelCase      :  Any						, )		->     int:
											A      						=    vocab_size
											A      						=    []
											self.cutoffs.extend(__UpperCamelCase       )
											if proj_share_all_but_first:
																		A      						=    [False] + [True] * len(self.cutoffs       )
											else:
																		A      						=    [False] + [False] * len(self.cutoffs       )
											A      						=    d_model
											A      						=    d_embed
											A      						=    d_head
											A      						=    d_inner
											A      						=    div_val
											A      						=    pre_lnorm
											A      						=    n_layer
											A      						=    n_head
											A      						=    mem_len
											A      						=    same_length
											A      						=    attn_type
											A      						=    clamp_len
											A      						=    sample_softmax
											A      						=    adaptive
											A      						=    dropout
											A      						=    dropatt
											A      						=    untie_r
											A      						=    init
											A      						=    init_range
											A      						=    proj_init_std
											A      						=    init_std
											A      						=    layer_norm_epsilon
											super().__init__(eos_token_id=__UpperCamelCase						, **__UpperCamelCase       )


				@property
				def __UpperCamelCase						(							self      :  Any       )		->     Optional[int]:
											# Message copied from Transformer-XL documentation
											logger.info(f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''       )
											return -1


				@max_position_embeddings.setter
				def __UpperCamelCase						(							self      :  List[Any]						, __UpperCamelCase      :  List[Any]       )		->     List[Any]:
											# Message copied from Transformer-XL documentation
											raise NotImplementedError(
											    f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''       )