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







     def      lowercase	(							self			:			Union[str, Any]    ,      _lowerCamelCase			:			Tuple			):
       return self.lineara(self.batchnorm(self.lineara(__UpperCAmelCase			)			)			)






class       lowerCAmelCase__				(						__A	):







     def      lowercase	(							self			:			Any    ,      _lowerCamelCase			:			Optional[int]    ,      *_lowerCamelCase			:			Union[str, Any]    ,      **_lowerCamelCase			:			Any			):
       return (args[0] + 1,) + args[1:], kwargs






class       lowerCAmelCase__				(						__A	):







     def      lowercase	(							self			:			Tuple    ,      _lowerCamelCase			:			Tuple    ,      _lowerCamelCase			:			List[Any]			):
       return output + 1






class       lowerCAmelCase__				(						unittest.TestCase	):







     def      lowercase	(							self			:			Any			):
       _snake_case					     =       ModelForTest()
       _snake_case					     =       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      lowercase	(							self			:			Optional[int]			):
       _snake_case					     =       ModelForTest()
       _snake_case					     =       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      lowercase	(							self			:			Union[str, Any]			):
       _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(__UpperCAmelCase    ,      __UpperCAmelCase			)
       _snake_case					     =       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
       _snake_case					     =       PreForwardHook()
       add_hook_to_module(__UpperCAmelCase    ,      __UpperCAmelCase			)
       _snake_case					     =       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
       _snake_case					     =       SequentialHook(PreForwardHook()    ,      PreForwardHook()			)
       add_hook_to_module(__UpperCAmelCase    ,      __UpperCAmelCase			)

       _snake_case					     =       test_model(__UpperCAmelCase			)
       assert torch.allclose(__UpperCAmelCase    ,      __UpperCAmelCase    ,      atol=1e-5			)







     def      lowercase	(							self			:			Dict			):
       _snake_case					     =       ModelForTest()
       _snake_case					     =       torch.randn(2    ,      3			)
       _snake_case					     =       test_model(__UpperCAmelCase			)

       _snake_case					     =       PostForwardHook()
       add_hook_to_module(__UpperCAmelCase    ,      __UpperCAmelCase			)
       _snake_case					     =       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
       _snake_case					     =       PostForwardHook()
       add_hook_to_module(__UpperCAmelCase    ,      __UpperCAmelCase			)
       _snake_case					     =       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
       _snake_case					     =       SequentialHook(PostForwardHook()    ,      PostForwardHook()			)
       add_hook_to_module(__UpperCAmelCase    ,      __UpperCAmelCase			)

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







     def      lowercase	(							self			:			List[Any]			):
       _snake_case					     =       ModelForTest()
       _snake_case					     =       torch.randn(2    ,      3			)
       _snake_case					     =       test_model(__UpperCAmelCase			)

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

       _snake_case					     =       True
       _snake_case					     =       test_model(__UpperCAmelCase			)
       self.assertFalse(outputa.requires_grad			)







     @require_multi_gpu
     def      lowercase	(							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
       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(__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			)			)
       _snake_case					     =       torch.randn(2    ,      3			).to(0			)
       _snake_case					     =       model(__UpperCAmelCase			)
       self.assertEqual(output.device    ,      torch.device(0			)			)







     def      lowercase	(							self			:			Tuple			):
       _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(**__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
       _snake_case					     =       torch.device(hook_kwargs['''execution_device''']			)
       self.assertEqual(model.batchnorm.running_mean.device    ,      __UpperCAmelCase			)

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

       _snake_case					     =       torch.randn(2    ,      3			)
       _snake_case					     =       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      lowercase	(							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					     =       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
       _snake_case					     =       torch.device(__UpperCAmelCase			)
       self.assertEqual(model.batchnorm.running_mean.device    ,      __UpperCAmelCase			)

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

       _snake_case					     =       torch.randn(2    ,      3			)
       _snake_case					     =       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      lowercase	(							self			:			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(
           __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
       _snake_case					     =       torch.device(__UpperCAmelCase			)
       self.assertEqual(model.batchnorm.running_mean.device    ,      __UpperCAmelCase			)

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

       _snake_case					     =       torch.randn(2    ,      3			)
       _snake_case					     =       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 List, Optional, Union

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


__lowerCamelCase							 =     logging.get_logger(__name__)

__lowerCamelCase							 =     {
    "huggingface/time-series-transformer-tourism-monthly": (
        "https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}
class 			UpperCamelCase__(	__A       ):
						lowerCAmelCase__     :  List[Any]  							=						'time_series_transformer'
						lowerCAmelCase__     :  Dict  							=						{
						    'hidden_size': 'd_model',
						    'num_attention_heads': 'encoder_attention_heads',
						    'num_hidden_layers': 'encoder_layers',
						}





						def __init__(				self						,__UpperCAmelCase = None						,__UpperCAmelCase = None						,__UpperCAmelCase = "student_t"						,__UpperCAmelCase = "nll"						,__UpperCAmelCase = 1						,__UpperCAmelCase = [1, 2, 3, 4, 5, 6, 7]						,__UpperCAmelCase = "mean"						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = 0						,__UpperCAmelCase = None						,__UpperCAmelCase = None						,__UpperCAmelCase = 32						,__UpperCAmelCase = 32						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = 2						,__UpperCAmelCase = True						,__UpperCAmelCase = "gelu"						,__UpperCAmelCase = 64						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 0.1						,__UpperCAmelCase = 1_00						,__UpperCAmelCase = 0.0_2						,__UpperCAmelCase=True						,**__UpperCAmelCase						,)    -> List[str]:
													# time series specific configuration
													A__		   =							prediction_length
													A__		   =							context_length or prediction_length
													A__		   =							distribution_output
													A__		   =							loss
													A__		   =							input_size
													A__		   =							num_time_features
													A__		   =							lags_sequence
													A__		   =							scaling
													A__		   =							num_dynamic_real_features
													A__		   =							num_static_real_features
													A__		   =							num_static_categorical_features
													if cardinality and num_static_categorical_features > 0:
																				if len(__UpperCAmelCase   ) != num_static_categorical_features:
																											raise ValueError(
																											    'The cardinality should be a list of the same length as `num_static_categorical_features`'   )
																				A__		   =							cardinality
													else:
																				A__		   =							[0]
													if embedding_dimension and num_static_categorical_features > 0:
																				if len(__UpperCAmelCase   ) != num_static_categorical_features:
																											raise ValueError(
																											    'The embedding dimension should be a list of the same length as `num_static_categorical_features`'   )
																				A__		   =							embedding_dimension
													else:
																				A__		   =							[min(50						,(cat + 1) // 2   ) for cat in self.cardinality]
													A__		   =							num_parallel_samples

													# Transformer architecture configuration
													A__		   =							input_size * len(__UpperCAmelCase   ) + self._number_of_features
													A__		   =							d_model
													A__		   =							encoder_attention_heads
													A__		   =							decoder_attention_heads
													A__		   =							encoder_ffn_dim
													A__		   =							decoder_ffn_dim
													A__		   =							encoder_layers
													A__		   =							decoder_layers

													A__		   =							dropout
													A__		   =							attention_dropout
													A__		   =							activation_dropout
													A__		   =							encoder_layerdrop
													A__		   =							decoder_layerdrop

													A__		   =							activation_function
													A__		   =							init_std

													A__		   =							use_cache

													super().__init__(is_encoder_decoder=__UpperCAmelCase						,**__UpperCAmelCase   )





						@property
						def  snake_case__		(				self   )    -> int:
													return (
													    sum(self.embedding_dimension   )
													    + self.num_dynamic_real_features
													    + self.num_time_features
													    + self.num_static_real_features
													    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
													)





