





'''simple docstring'''



import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

from transformers import is_torch_available
from transformers.testing_utils import is_flaky, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


_SCREAMING_SNAKE_CASE   =      1e-4

if is_torch_available():
	import torch

	from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
	from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder






@require_torch
class 							__lowercase    :







	'''simple docstring'''

	def __init__(self       ,_lowerCamelCase       ,_lowerCamelCase=16       ,_lowerCamelCase=13       ,_lowerCamelCase=7       ,_lowerCamelCase=14       ,_lowerCamelCase=10       ,_lowerCamelCase=19       ,_lowerCamelCase=5       ,_lowerCamelCase=4       ,_lowerCamelCase=True       ,_lowerCamelCase=16       ,_lowerCamelCase=2       ,_lowerCamelCase=4       ,_lowerCamelCase=4       ,_lowerCamelCase="gelu"       ,_lowerCamelCase=0.1       ,_lowerCamelCase=0.1       ,_lowerCamelCase=[1, 2, 3, 4, 5]       ,_lowerCamelCase=25       ,_lowerCamelCase=5       ,)							->							List[Any]:







				'''simple docstring'''





				__lowercase			  =							d_model
				__lowercase			  =							parent
				__lowercase			  =							batch_size
				__lowercase			  =							prediction_length
				__lowercase			  =							context_length
				__lowercase			  =							cardinality
				__lowercase			  =							num_time_features
				__lowercase			  =							lags_sequence
				__lowercase			  =							embedding_dimension
				__lowercase			  =							is_training
				__lowercase			  =							hidden_size
				__lowercase			  =							num_hidden_layers
				__lowercase			  =							num_attention_heads
				__lowercase			  =							intermediate_size
				__lowercase			  =							hidden_act
				__lowercase			  =							hidden_dropout_prob
				__lowercase			  =							attention_probs_dropout_prob

				__lowercase			  =							context_length
				__lowercase			  =							prediction_length + label_length
				__lowercase			  =							label_length

				__lowercase			  =							moving_average
				__lowercase			  =							autocorrelation_factor


	def   _UpperCAmelCase       (self	)							->							List[str]:







				'''simple docstring'''





				return AutoformerConfig(
				    d_model=self.d_model       ,encoder_layers=self.num_hidden_layers       ,decoder_layers=self.num_hidden_layers       ,encoder_attention_heads=self.num_attention_heads       ,decoder_attention_heads=self.num_attention_heads       ,encoder_ffn_dim=self.intermediate_size       ,decoder_ffn_dim=self.intermediate_size       ,dropout=self.hidden_dropout_prob       ,attention_dropout=self.attention_probs_dropout_prob       ,prediction_length=self.prediction_length       ,context_length=self.context_length       ,label_length=self.label_length       ,lags_sequence=self.lags_sequence       ,num_time_features=self.num_time_features       ,num_static_categorical_features=1       ,cardinality=[self.cardinality]       ,embedding_dimension=[self.embedding_dimension]       ,moving_average=self.moving_average       ,)


	def   _UpperCAmelCase       (self       ,_lowerCamelCase	)							->							List[Any]:







				'''simple docstring'''





				__lowercase			  =							config.context_length + max(config.lags_sequence	)

				__lowercase			  =							ids_tensor([self.batch_size, 1]       ,config.cardinality[0]	)
				__lowercase			  =							floats_tensor([self.batch_size, _past_length, config.num_time_features]	)
				__lowercase			  =							floats_tensor([self.batch_size, _past_length]	)
				__lowercase			  =							floats_tensor([self.batch_size, _past_length]	) > 0.5

				# decoder inputs
				__lowercase			  =							floats_tensor([self.batch_size, config.prediction_length, config.num_time_features]	)
				__lowercase			  =							floats_tensor([self.batch_size, config.prediction_length]	)

				__lowercase			  =							{
				    '''past_values''': past_values,
				    '''static_categorical_features''': static_categorical_features,
				    '''past_time_features''': past_time_features,
				    '''past_observed_mask''': past_observed_mask,
				    '''future_time_features''': future_time_features,
				    '''future_values''': future_values,
				}
				return inputs_dict


	def   _UpperCAmelCase       (self	)							->							int:







				'''simple docstring'''





				__lowercase			  =							self.get_config()
				__lowercase			  =							self.prepare_autoformer_inputs_dict(_lowerCamelCase	)
				return config, inputs_dict


	def   _UpperCAmelCase       (self	)							->							Dict:







				'''simple docstring'''





				__lowercase  ,						__lowercase			  =							self.prepare_config_and_inputs()
				return config, inputs_dict


	def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase	)							->							str:







				'''simple docstring'''





				__lowercase			  =							AutoformerModel(config=_lowerCamelCase	).to(_lowerCamelCase	).eval()
				__lowercase			  =							model(**_lowerCamelCase	)

				__lowercase			  =							outputs.encoder_last_hidden_state
				__lowercase			  =							outputs.last_hidden_state

				with tempfile.TemporaryDirectory() as tmpdirname:
							__lowercase			  =							model.get_encoder()
							encoder.save_pretrained(_lowerCamelCase	)
							__lowercase			  =							AutoformerEncoder.from_pretrained(_lowerCamelCase	).to(_lowerCamelCase	)

				__lowercase  ,						__lowercase  ,						__lowercase  ,						__lowercase  ,						__lowercase			  =							model.create_network_inputs(**_lowerCamelCase	)
				__lowercase  ,						__lowercase			  =							model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]	)

				__lowercase			  =							torch.cat(
				    (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...])       ,dim=-1       ,)
				__lowercase			  =							encoder(inputs_embeds=_lowerCamelCase	)[0]
				self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1E-3	)

				__lowercase			  =							(
				    torch.mean(transformer_inputs[:, : config.context_length, ...]       ,dim=1	)
				    .unsqueeze(1	)
				    .repeat(1       ,config.prediction_length       ,1	)
				)
				__lowercase			  =							torch.zeros(
				    [transformer_inputs.shape[0], config.prediction_length, transformer_inputs.shape[2]]       ,device=enc_input.device       ,)

				__lowercase			  =							torch.cat(
				    (
				        torch.cat((seasonal_input[:, -config.label_length :, ...], zeros)       ,dim=1	),
				        feature[:, config.context_length - config.label_length :, ...],
				    )       ,dim=-1       ,)
				__lowercase			  =							torch.cat(
				    (
				        torch.cat((trend_input[:, -config.label_length :, ...], mean)       ,dim=1	),
				        feature[:, config.context_length - config.label_length :, ...],
				    )       ,dim=-1       ,)

				with tempfile.TemporaryDirectory() as tmpdirname:
							__lowercase			  =							model.get_decoder()
							decoder.save_pretrained(_lowerCamelCase	)
							__lowercase			  =							AutoformerDecoder.from_pretrained(_lowerCamelCase	).to(_lowerCamelCase	)

				__lowercase			  =							decoder(
				    trend=_lowerCamelCase       ,inputs_embeds=_lowerCamelCase       ,encoder_hidden_states=_lowerCamelCase       ,)[0]

				self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1E-3	)






@require_torch
class 							__lowercase    ( lowerCAmelCase__	, lowerCAmelCase__	, unittest.TestCase  ):







	'''simple docstring'''

	a      :			Optional[Any]     =    (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
	a      :			List[str]     =    (AutoformerForPrediction,) if is_torch_available() else ()
	a      :			Optional[Any]     =    {"feature-extraction": AutoformerModel} if is_torch_available() else {}
	a      :			Optional[int]     =    False
	a      :			Tuple     =    False
	a      :			Optional[int]     =    False
	a      :			Optional[Any]     =    False
	a      :			Optional[int]     =    False
	a      :			Optional[int]     =    False


	def   _UpperCAmelCase       (self	)							->							Dict:







				'''simple docstring'''





				__lowercase			  =							AutoformerModelTester(self	)
				__lowercase			  =							ConfigTester(self       ,config_class=_lowerCamelCase       ,has_text_modality=_lowerCamelCase	)


	def   _UpperCAmelCase       (self	)							->							List[Any]:







				'''simple docstring'''





				self.config_tester.run_common_tests()


	def   _UpperCAmelCase       (self	)							->							Union[str, Any]:







				'''simple docstring'''





				__lowercase  ,						__lowercase			  =							self.model_tester.prepare_config_and_inputs()
				for model_class in self.all_model_classes:
							__lowercase			  =							model_class(_lowerCamelCase	)

							with tempfile.TemporaryDirectory() as tmpdirname:
										model.save_pretrained(_lowerCamelCase	)
										__lowercase  ,						__lowercase			  =							model_class.from_pretrained(_lowerCamelCase       ,output_loading_info=_lowerCamelCase	)
							self.assertEqual(info['''missing_keys''']       ,[]	)


	def   _UpperCAmelCase       (self	)							->							List[str]:







				'''simple docstring'''





				__lowercase			  =							self.model_tester.prepare_config_and_inputs_for_common()
				self.model_tester.check_encoder_decoder_model_standalone(*_lowerCamelCase	)


	@unittest.skip(reason='''Model has no tokens embeddings'''	)
	def   _UpperCAmelCase       (self	)							->							Optional[Any]:







				'''simple docstring'''





				pass


	def   _UpperCAmelCase       (self	)							->							Tuple:







				'''simple docstring'''





				__lowercase			  =							inspect.signature(getattr(_lowerCamelCase       ,'''forward'''	)	)
				# The main input is the name of the argument after `self`
				__lowercase			  =							list(model_signature.parameters.keys()	)[1]
				self.assertEqual(AutoformerModel.main_input_name       ,_lowerCamelCase	)


	def   _UpperCAmelCase       (self	)							->							Dict:







				'''simple docstring'''





				__lowercase  ,						__lowercase			  =							self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
							__lowercase			  =							model_class(_lowerCamelCase	)
							__lowercase			  =							inspect.signature(model.forward	)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							__lowercase			  =							[*signature.parameters.keys()]

							__lowercase			  =							[
							    '''past_values''',
							    '''past_time_features''',
							    '''past_observed_mask''',
							    '''static_categorical_features''',
							    '''static_real_features''',
							    '''future_values''',
							    '''future_time_features''',
							]

							if model.__class__.__name__ in ["AutoformerForPrediction"]:
										expected_arg_names.append('''future_observed_mask'''	)

							expected_arg_names.extend(
							    [
							        '''decoder_attention_mask''',
							        '''head_mask''',
							        '''decoder_head_mask''',
							        '''cross_attn_head_mask''',
							        '''encoder_outputs''',
							        '''past_key_values''',
							        '''output_hidden_states''',
							        '''output_attentions''',
							        '''use_cache''',
							        '''return_dict''',
							    ]	)

							self.assertListEqual(arg_names[: len(_lowerCamelCase	)]       ,_lowerCamelCase	)


	def   _UpperCAmelCase       (self	)							->							Optional[Any]:







				'''simple docstring'''





				__lowercase  ,						__lowercase			  =							self.model_tester.prepare_config_and_inputs_for_common()
				__lowercase			  =							True

				__lowercase			  =							getattr(self.model_tester       ,'''seq_length'''       ,_lowerCamelCase	)
				__lowercase			  =							getattr(self.model_tester       ,'''decoder_seq_length'''       ,_lowerCamelCase	)
				__lowercase			  =							getattr(self.model_tester       ,'''encoder_seq_length'''       ,_lowerCamelCase	)
				__lowercase			  =							getattr(self.model_tester       ,'''d_model'''       ,_lowerCamelCase	)
				__lowercase			  =							getattr(self.model_tester       ,'''num_attention_heads'''       ,_lowerCamelCase	)
				__lowercase			  =							d_model // num_attention_heads

				for model_class in self.all_model_classes:
							__lowercase			  =							True
							__lowercase			  =							False
							__lowercase			  =							True
							__lowercase			  =							model_class(_lowerCamelCase	)
							model.to(_lowerCamelCase	)
							model.eval()
							with torch.no_grad():
										__lowercase			  =							model(**self._prepare_for_class(_lowerCamelCase       ,_lowerCamelCase	)	)
							__lowercase			  =							outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
							self.assertEqual(len(_lowerCamelCase	)       ,self.model_tester.num_hidden_layers	)

							# check that output_attentions also work using config
							del inputs_dict["output_attentions"]
							__lowercase			  =							True
							__lowercase			  =							model_class(_lowerCamelCase	)
							model.to(_lowerCamelCase	)
							model.eval()
							with torch.no_grad():
										__lowercase			  =							model(**self._prepare_for_class(_lowerCamelCase       ,_lowerCamelCase	)	)
							__lowercase			  =							outputs.encoder_attentions
							self.assertEqual(len(_lowerCamelCase	)       ,self.model_tester.num_hidden_layers	)

							self.assertListEqual(
							    list(attentions[0].shape[-3:]	)       ,[self.model_tester.num_attention_heads, encoder_seq_length, dim]       ,)
							__lowercase			  =							len(_lowerCamelCase	)

							__lowercase			  =							7

							if "last_hidden_state" in outputs:
										correct_outlen += 1

							if "trend" in outputs:
										correct_outlen += 1

							if "past_key_values" in outputs:
										correct_outlen += 1  # past_key_values have been returned

							if "loss" in outputs:
										correct_outlen += 1

							if "params" in outputs:
										correct_outlen += 1

							self.assertEqual(_lowerCamelCase       ,_lowerCamelCase	)

							# decoder attentions
							__lowercase			  =							outputs.decoder_attentions
							self.assertIsInstance(_lowerCamelCase       ,(list, tuple)	)
							self.assertEqual(len(_lowerCamelCase	)       ,self.model_tester.num_hidden_layers	)
							self.assertListEqual(
							    list(decoder_attentions[0].shape[-3:]	)       ,[self.model_tester.num_attention_heads, decoder_seq_length, dim]       ,)

							# cross attentions
							__lowercase			  =							outputs.cross_attentions
							self.assertIsInstance(_lowerCamelCase       ,(list, tuple)	)
							self.assertEqual(len(_lowerCamelCase	)       ,self.model_tester.num_hidden_layers	)
							self.assertListEqual(
							    list(cross_attentions[0].shape[-3:]	)       ,[self.model_tester.num_attention_heads, decoder_seq_length, dim]       ,)

				# Check attention is always last and order is fine
				__lowercase			  =							True
				__lowercase			  =							True
				__lowercase			  =							model_class(_lowerCamelCase	)
				model.to(_lowerCamelCase	)
				model.eval()
				with torch.no_grad():
							__lowercase			  =							model(**self._prepare_for_class(_lowerCamelCase       ,_lowerCamelCase	)	)

				self.assertEqual(out_len + 2       ,len(_lowerCamelCase	)	)

				__lowercase			  =							outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

				self.assertEqual(len(_lowerCamelCase	)       ,self.model_tester.num_hidden_layers	)
				self.assertListEqual(
				    list(self_attentions[0].shape[-3:]	)       ,[self.model_tester.num_attention_heads, encoder_seq_length, dim]       ,)


	@is_flaky()
	def   _UpperCAmelCase       (self	)							->							Dict:







				'''simple docstring'''





				super().test_retain_grad_hidden_states_attentions()





def 				_lowerCAmelCase  (			lowerCamelCase_     :					Any="train-batch.pt"			):
			__lowercase			  =							hf_hub_download(repo_id='''hf-internal-testing/tourism-monthly-batch'''						,    filename=lowerCamelCase_						,    repo_type='''dataset'''			)
			__lowercase			  =							torch.load(lowerCamelCase_						,    map_location=lowerCamelCase_			)
			return batch






@require_torch
@slow
class 							__lowercase    ( unittest.TestCase  ):







	'''simple docstring'''

	def   _UpperCAmelCase       (self	)							->							Union[str, Any]:







				'''simple docstring'''





				__lowercase			  =							AutoformerModel.from_pretrained('''huggingface/autoformer-tourism-monthly'''	).to(_lowerCamelCase	)
				__lowercase			  =							prepare_batch()

				with torch.no_grad():
							__lowercase			  =							model(
							    past_values=batch['''past_values''']       ,past_time_features=batch['''past_time_features''']       ,past_observed_mask=batch['''past_observed_mask''']       ,static_categorical_features=batch['''static_categorical_features''']       ,future_values=batch['''future_values''']       ,future_time_features=batch['''future_time_features''']       ,)[0]

				__lowercase			  =							torch.Size(
				    (64, model.config.prediction_length + model.config.label_length, model.config.feature_size)	)
				self.assertEqual(output.shape       ,_lowerCamelCase	)

				__lowercase			  =							torch.tensor(
				    [[0.3_5_9_3, -1.3_3_9_8, 0.6_3_3_0], [0.2_2_7_9, 1.5_3_9_6, -0.1_7_9_2], [0.0_4_5_0, 1.3_2_2_5, -0.2_3_3_5]]       ,device=_lowerCamelCase	)
				self.assertTrue(torch.allclose(output[0, :3, :3]       ,_lowerCamelCase       ,atol=_lowerCamelCase	)	)


	def   _UpperCAmelCase       (self	)							->							Any:







				'''simple docstring'''





				__lowercase			  =							AutoformerForPrediction.from_pretrained('''huggingface/autoformer-tourism-monthly'''	).to(_lowerCamelCase	)
				__lowercase			  =							prepare_batch('''val-batch.pt'''	)
				with torch.no_grad():
							__lowercase			  =							model(
							    past_values=batch['''past_values''']       ,past_time_features=batch['''past_time_features''']       ,past_observed_mask=batch['''past_observed_mask''']       ,static_categorical_features=batch['''static_categorical_features''']       ,).encoder_last_hidden_state
				__lowercase			  =							torch.Size((64, model.config.context_length, model.config.d_model)	)
				self.assertEqual(output.shape       ,_lowerCamelCase	)

				__lowercase			  =							torch.tensor(
				    [[-0.0_7_3_4, -0.9_0_3_6, 0.8_3_5_8], [4.7_1_8_6, 2.4_1_1_3, 1.9_5_8_1], [1.7_9_5_3, 2.3_5_5_8, 1.2_9_7_0]]       ,device=_lowerCamelCase	)
				self.assertTrue(torch.allclose(output[0, :3, :3]       ,_lowerCamelCase       ,atol=_lowerCamelCase	)	)


	def   _UpperCAmelCase       (self	)							->							int:







				'''simple docstring'''





				__lowercase			  =							AutoformerForPrediction.from_pretrained('''huggingface/autoformer-tourism-monthly'''	).to(_lowerCamelCase	)
				__lowercase			  =							prepare_batch('''val-batch.pt'''	)
				with torch.no_grad():
							__lowercase			  =							model.generate(
							    static_categorical_features=batch['''static_categorical_features''']       ,past_time_features=batch['''past_time_features''']       ,past_values=batch['''past_values''']       ,future_time_features=batch['''future_time_features''']       ,past_observed_mask=batch['''past_observed_mask''']       ,)
				__lowercase			  =							torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length)	)
				self.assertEqual(outputs.sequences.shape       ,_lowerCamelCase	)

				__lowercase			  =							torch.tensor([3_1_3_0.6_7_6_3, 4_0_5_6.5_2_9_3, 7_0_5_3.0_7_8_6]       ,device=_lowerCamelCase	)
				__lowercase			  =							outputs.sequences.mean(dim=1	)
				self.assertTrue(torch.allclose(mean_prediction[0, -3:]       ,_lowerCamelCase       ,rtol=1E-1	)	)












'''simple docstring'''



def 				_lowerCAmelCase  (			lowerCamelCase_     :					Union[str, Any]			):
			__lowercase			  =							1
			__lowercase			  =							2
			while i * i <= n:
						__lowercase			  =							0
						while n % i == 0:
									n //= i
									multiplicity += 1
						n_divisors *= multiplicity + 1
						i += 1
			if n > 1:
						n_divisors *= 2
			return n_divisors




def 				_lowerCAmelCase  (			):
			__lowercase			  =							1
			__lowercase			  =							1

			while True:
						i += 1
						t_num += i

						if count_divisors(lowerCamelCase_			) > 5_0_0:
									break

			return t_num


if __name__ == "__main__":
	print(solution())





