


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


__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    UpperCAmelCase					:




			def __init__(						self							:		str       ,       __magic_name__							:		List[str]       ,       __magic_name__							:		Tuple=1_6       ,       __magic_name__							:		Dict=1_3       ,       __magic_name__							:		List[Any]=7       ,       __magic_name__							:		List[str]=1_4       ,       __magic_name__							:		Optional[Any]=1_0       ,       __magic_name__							:		List[str]=1_9       ,       __magic_name__							:		str=5       ,       __magic_name__							:		Optional[Any]=4       ,       __magic_name__							:		Optional[Any]=True       ,       __magic_name__							:		List[Any]=1_6       ,       __magic_name__							:		Optional[int]=2       ,       __magic_name__							:		List[str]=4       ,       __magic_name__							:		Dict=4       ,       __magic_name__							:		List[str]="gelu"       ,       __magic_name__							:		List[Any]=0.1       ,       __magic_name__							:		Tuple=0.1       ,       __magic_name__							:		int=[1, 2, 3, 4, 5]       ,       __magic_name__							:		Tuple=2_5       ,       __magic_name__							:		List[Any]=5       ,       ):







						"""simple docstring"""







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

						UpperCamelCase         =					context_length
						UpperCamelCase         =					prediction_length + label_length
						UpperCamelCase         =					label_length

						UpperCamelCase         =					moving_average
						UpperCamelCase         =					autocorrelation_factor




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







						"""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   lowerCamelCase_					(						self							:		Dict       ,       __magic_name__							:		Dict						):







						"""simple docstring"""







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

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

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

						UpperCamelCase         =					{
						    """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   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						UpperCamelCase         =					self.get_config()
						UpperCamelCase         =					self.prepare_autoformer_inputs_dict(__magic_name__						)
						return config, inputs_dict




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







						"""simple docstring"""







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





			def   lowerCamelCase_					(						self							:		Any       ,       __magic_name__							:		Optional[Any]       ,       __magic_name__							:		Any						):







						"""simple docstring"""







						UpperCamelCase         =					AutoformerModel(config=__magic_name__						).to(__magic_name__						).eval()
						UpperCamelCase         =					model(**__magic_name__						)

						UpperCamelCase         =					outputs.encoder_last_hidden_state
						UpperCamelCase         =					outputs.last_hidden_state

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

						UpperCamelCase  ,  UpperCamelCase  ,  UpperCamelCase  ,  UpperCamelCase  ,  UpperCamelCase         =					model.create_network_inputs(**__magic_name__						)
						UpperCamelCase  ,  UpperCamelCase         =					model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]						)

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

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

						UpperCamelCase         =					torch.cat(
						    (
						        torch.cat((seasonal_input[:, -config.label_length :, ...], zeros)       ,       dim=1						),
						        feature[:, config.context_length - config.label_length :, ...],
						    )       ,       dim=-1       ,       )
						UpperCamelCase         =					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:
									UpperCamelCase         =					model.get_decoder()
									decoder.save_pretrained(__magic_name__						)
									UpperCamelCase         =					AutoformerDecoder.from_pretrained(__magic_name__						).to(__magic_name__						)

						UpperCamelCase         =					decoder(
						    trend=__magic_name__       ,       inputs_embeds=__magic_name__       ,       encoder_hidden_states=__magic_name__       ,       )[0]

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




@require_torch
class    UpperCAmelCase					(       __snake_case      , __snake_case      , unittest.TestCase    ):
			lowercase						= (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
			lowercase						= (AutoformerForPrediction,) if is_torch_available() else ()
			lowercase						= {"""feature-extraction""": AutoformerModel} if is_torch_available() else {}
			lowercase						= False
			lowercase						= False
			lowercase						= False
			lowercase						= False
			lowercase						= False
			lowercase						= False




			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						UpperCamelCase         =					AutoformerModelTester(self						)
						UpperCamelCase         =					ConfigTester(self       ,       config_class=__magic_name__       ,       has_text_modality=__magic_name__						)




			def   lowerCamelCase_					(						self							:		int						):







						"""simple docstring"""







						self.config_tester.run_common_tests()




			def   lowerCamelCase_					(						self							:		List[str]						):







						"""simple docstring"""







						UpperCamelCase  ,  UpperCamelCase         =					self.model_tester.prepare_config_and_inputs()
						for model_class in self.all_model_classes:
									UpperCamelCase         =					model_class(__magic_name__						)

									with tempfile.TemporaryDirectory() as tmpdirname:
												model.save_pretrained(__magic_name__						)
												UpperCamelCase  ,  UpperCamelCase         =					model_class.from_pretrained(__magic_name__       ,       output_loading_info=__magic_name__						)
									self.assertEqual(info["""missing_keys"""]       ,       []						)




			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						UpperCamelCase         =					self.model_tester.prepare_config_and_inputs_for_common()
						self.model_tester.check_encoder_decoder_model_standalone(*__magic_name__						)




			@unittest.skip(reason="""Model has no tokens embeddings"""						)
			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						pass




			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







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




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







						"""simple docstring"""







						UpperCamelCase  ,  UpperCamelCase         =					self.model_tester.prepare_config_and_inputs_for_common()

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

									UpperCamelCase         =					[
									    """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(__magic_name__						)]       ,       __magic_name__						)




			def   lowerCamelCase_					(						self							:		Optional[int]						):







						"""simple docstring"""







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

						UpperCamelCase         =					getattr(self.model_tester       ,       """seq_length"""       ,       __magic_name__						)
						UpperCamelCase         =					getattr(self.model_tester       ,       """decoder_seq_length"""       ,       __magic_name__						)
						UpperCamelCase         =					getattr(self.model_tester       ,       """encoder_seq_length"""       ,       __magic_name__						)
						UpperCamelCase         =					getattr(self.model_tester       ,       """d_model"""       ,       __magic_name__						)
						UpperCamelCase         =					getattr(self.model_tester       ,       """num_attention_heads"""       ,       __magic_name__						)
						UpperCamelCase         =					d_model // num_attention_heads

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

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

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

									UpperCamelCase         =					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(__magic_name__       ,       __magic_name__						)

									# decoder attentions
									UpperCamelCase         =					outputs.decoder_attentions
									self.assertIsInstance(__magic_name__       ,       (list, tuple)						)
									self.assertEqual(len(__magic_name__						)       ,       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
									UpperCamelCase         =					outputs.cross_attentions
									self.assertIsInstance(__magic_name__       ,       (list, tuple)						)
									self.assertEqual(len(__magic_name__						)       ,       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
						UpperCamelCase         =					True
						UpperCamelCase         =					True
						UpperCamelCase         =					model_class(__magic_name__						)
						model.to(__magic_name__						)
						model.eval()
						with torch.no_grad():
									UpperCamelCase         =					model(**self._prepare_for_class(__magic_name__       ,       __magic_name__						)						)

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

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

						self.assertEqual(len(__magic_name__						)       ,       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   lowerCamelCase_					(						self							:		int						):







						"""simple docstring"""







						super().test_retain_grad_hidden_states_attentions()






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int="train-batch.pt"	):





			"""simple docstring"""




			UpperCamelCase         =					hf_hub_download(repo_id="""hf-internal-testing/tourism-monthly-batch"""       ,  filename=SCREAMING_SNAKE_CASE_       ,  repo_type="""dataset"""	)
			UpperCamelCase         =					torch.load(SCREAMING_SNAKE_CASE_       ,  map_location=SCREAMING_SNAKE_CASE_	)
			return batch




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




			def   lowerCamelCase_					(						self							:		Dict						):







						"""simple docstring"""







						UpperCamelCase         =					AutoformerModel.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(__magic_name__						)
						UpperCamelCase         =					prepare_batch()

						with torch.no_grad():
									UpperCamelCase         =					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]

						UpperCamelCase         =					torch.Size(
						    (6_4, model.config.prediction_length + model.config.label_length, model.config.feature_size)						)
						self.assertEqual(output.shape       ,       __magic_name__						)

						UpperCamelCase         =					torch.tensor(
						    [[0.3_593, -1.3_398, 0.6_330], [0.2_279, 1.5_396, -0.1_792], [0.0_450, 1.3_225, -0.2_335]]       ,       device=__magic_name__						)
						self.assertTrue(torch.allclose(output[0, :3, :3]       ,       __magic_name__       ,       atol=__magic_name__						)						)




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







						"""simple docstring"""







						UpperCamelCase         =					AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(__magic_name__						)
						UpperCamelCase         =					prepare_batch("""val-batch.pt"""						)
						with torch.no_grad():
									UpperCamelCase         =					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
						UpperCamelCase         =					torch.Size((6_4, model.config.context_length, model.config.d_model)						)
						self.assertEqual(output.shape       ,       __magic_name__						)

						UpperCamelCase         =					torch.tensor(
						    [[-0.0_734, -0.9_036, 0.8_358], [4.7_186, 2.4_113, 1.9_581], [1.7_953, 2.3_558, 1.2_970]]       ,       device=__magic_name__						)
						self.assertTrue(torch.allclose(output[0, :3, :3]       ,       __magic_name__       ,       atol=__magic_name__						)						)





			def   lowerCamelCase_					(						self							:		Any						):







						"""simple docstring"""







						UpperCamelCase         =					AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(__magic_name__						)
						UpperCamelCase         =					prepare_batch("""val-batch.pt"""						)
						with torch.no_grad():
									UpperCamelCase         =					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"""]       ,       )
						UpperCamelCase         =					torch.Size((6_4, model.config.num_parallel_samples, model.config.prediction_length)						)
						self.assertEqual(outputs.sequences.shape       ,       __magic_name__						)

						UpperCamelCase         =					torch.tensor([3_130.6_763, 4_056.5_293, 7_053.0_786]       ,       device=__magic_name__						)
						UpperCamelCase         =					outputs.sequences.mean(dim=1						)
						self.assertTrue(torch.allclose(mean_prediction[0, -3:]       ,       __magic_name__       ,       rtol=1e-1						)						)








import math
import unittest






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			assert isinstance(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	) and (
			    number >= 0
			), "'number' must been an int and positive"

			if 1 < number < 4:
						# 2 and 3 are primes
						return True
			elif number < 2 or number % 2 == 0 or number % 3 == 0:
						# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
						return False

			# All primes number are in format of 6k +/- 1
			for i in range(5       ,  int(math.sqrt(SCREAMING_SNAKE_CASE_	) + 1	)       ,  6	):
						if number % i == 0 or number % (i + 2) == 0:
									return False
			return True




class    UpperCAmelCase					(       unittest.TestCase    ):




			def   lowerCamelCase_					(						self							:		Any						):







						"""simple docstring"""







						self.assertTrue(is_prime(2						)						)
						self.assertTrue(is_prime(3						)						)
						self.assertTrue(is_prime(5						)						)
						self.assertTrue(is_prime(7						)						)
						self.assertTrue(is_prime(1_1						)						)
						self.assertTrue(is_prime(1_3						)						)
						self.assertTrue(is_prime(1_7						)						)
						self.assertTrue(is_prime(1_9						)						)
						self.assertTrue(is_prime(2_3						)						)
						self.assertTrue(is_prime(2_9						)						)





			def   lowerCamelCase_					(						self							:		Dict						):







						"""simple docstring"""







						with self.assertRaises(__magic_name__						):
									is_prime(-1_9						)
						self.assertFalse(
						    is_prime(0						)       ,       """Zero doesn't have any positive factors, primes must have exactly two."""       ,       )
						self.assertFalse(
						    is_prime(1						)       ,       """One only has 1 positive factor, primes must have exactly two."""       ,       )
						self.assertFalse(is_prime(2 * 2						)						)
						self.assertFalse(is_prime(2 * 3						)						)
						self.assertFalse(is_prime(3 * 3						)						)
						self.assertFalse(is_prime(3 * 5						)						)
						self.assertFalse(is_prime(3 * 5 * 7						)						)


if __name__ == "__main__":
			unittest.main()




