






"""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   :							Any          =		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 	a		:






    def __init__(			self	:				Dict			,		__SCREAMING_SNAKE_CASE	:				Union[str, Any]			,		__SCREAMING_SNAKE_CASE	:				List[str]=16			,		__SCREAMING_SNAKE_CASE	:				Optional[int]=13			,		__SCREAMING_SNAKE_CASE	:				Tuple=7			,		__SCREAMING_SNAKE_CASE	:				Tuple=14			,		__SCREAMING_SNAKE_CASE	:				int=10			,		__SCREAMING_SNAKE_CASE	:				List[Any]=19			,		__SCREAMING_SNAKE_CASE	:				Any=5			,		__SCREAMING_SNAKE_CASE	:				str=4			,		__SCREAMING_SNAKE_CASE	:				Tuple=True			,		__SCREAMING_SNAKE_CASE	:				List[Any]=16			,		__SCREAMING_SNAKE_CASE	:				Union[str, Any]=2			,		__SCREAMING_SNAKE_CASE	:				int=4			,		__SCREAMING_SNAKE_CASE	:				int=4			,		__SCREAMING_SNAKE_CASE	:				str="gelu"			,		__SCREAMING_SNAKE_CASE	:				Optional[Any]=0.1			,		__SCREAMING_SNAKE_CASE	:				Optional[int]=0.1			,		__SCREAMING_SNAKE_CASE	:				Optional[int]=[1, 2, 3, 4, 5]			,		__SCREAMING_SNAKE_CASE	:				List[Any]=25			,		__SCREAMING_SNAKE_CASE	:				Tuple=5			,		)	->						List[Any]:
      lowerCamelCase_       	=		d_model
      lowerCamelCase_       	=		parent
      lowerCamelCase_       	=		batch_size
      lowerCamelCase_       	=		prediction_length
      lowerCamelCase_       	=		context_length
      lowerCamelCase_       	=		cardinality
      lowerCamelCase_       	=		num_time_features
      lowerCamelCase_       	=		lags_sequence
      lowerCamelCase_       	=		embedding_dimension
      lowerCamelCase_       	=		is_training
      lowerCamelCase_       	=		hidden_size
      lowerCamelCase_       	=		num_hidden_layers
      lowerCamelCase_       	=		num_attention_heads
      lowerCamelCase_       	=		intermediate_size
      lowerCamelCase_       	=		hidden_act
      lowerCamelCase_       	=		hidden_dropout_prob
      lowerCamelCase_       	=		attention_probs_dropout_prob

      lowerCamelCase_       	=		context_length
      lowerCamelCase_       	=		prediction_length + label_length
      lowerCamelCase_       	=		label_length

      lowerCamelCase_       	=		moving_average
      lowerCamelCase_       	=		autocorrelation_factor






    def   UpperCamelCase			(			self	:				Tuple   )	->						List[str]:
      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	:				Tuple			,		__SCREAMING_SNAKE_CASE	:				Dict   )	->						Optional[Any]:
      lowerCamelCase_       	=		config.context_length + max(config.lags_sequence   )

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

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

      lowerCamelCase_       	=		{
          '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	:				Optional[int]   )	->						Optional[int]:
      lowerCamelCase_       	=		self.get_config()
      lowerCamelCase_       	=		self.prepare_autoformer_inputs_dict(__SCREAMING_SNAKE_CASE   )
      return config, inputs_dict






    def   UpperCamelCase			(			self	:				List[Any]   )	->						List[Any]:
      lowerCamelCase_	,      lowerCamelCase_       	=		self.prepare_config_and_inputs()
      return config, inputs_dict






    def   UpperCamelCase			(			self	:				Tuple			,		__SCREAMING_SNAKE_CASE	:				Dict			,		__SCREAMING_SNAKE_CASE	:				Tuple   )	->						int:
      lowerCamelCase_       	=		AutoformerModel(config=__SCREAMING_SNAKE_CASE   ).to(__SCREAMING_SNAKE_CASE   ).eval()
      lowerCamelCase_       	=		model(**__SCREAMING_SNAKE_CASE   )

      lowerCamelCase_       	=		outputs.encoder_last_hidden_state
      lowerCamelCase_       	=		outputs.last_hidden_state

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

      lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_       	=		model.create_network_inputs(**__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_	,      lowerCamelCase_       	=		model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]   )

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

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

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

      lowerCamelCase_       	=		decoder(
          trend=__SCREAMING_SNAKE_CASE			,		inputs_embeds=__SCREAMING_SNAKE_CASE			,		encoder_hidden_states=__SCREAMING_SNAKE_CASE			,		)[0]

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






@require_torch
class 	a		(   __snake_case    , __snake_case    , unittest.TestCase						):
    SCREAMING_SNAKE_CASE		:	Optional[int]		=					(AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
    SCREAMING_SNAKE_CASE		:	Union[str, Any]		=					(AutoformerForPrediction,) if is_torch_available() else ()
    SCREAMING_SNAKE_CASE		:	Optional[int]		=					{"""feature-extraction""": AutoformerModel} if is_torch_available() else {}
    SCREAMING_SNAKE_CASE		:	List[Any]		=					False
    SCREAMING_SNAKE_CASE		:	Union[str, Any]		=					False
    SCREAMING_SNAKE_CASE		:	List[str]		=					False
    SCREAMING_SNAKE_CASE		:	Optional[Any]		=					False
    SCREAMING_SNAKE_CASE		:	Tuple		=					False
    SCREAMING_SNAKE_CASE		:	str		=					False






    def   UpperCamelCase			(			self	:				Tuple   )	->						Dict:
      lowerCamelCase_       	=		AutoformerModelTester(self   )
      lowerCamelCase_       	=		ConfigTester(self			,		config_class=__SCREAMING_SNAKE_CASE			,		has_text_modality=__SCREAMING_SNAKE_CASE   )






    def   UpperCamelCase			(			self	:				Dict   )	->						Optional[int]:
      self.config_tester.run_common_tests()






    def   UpperCamelCase			(			self	:				int   )	->						int:
      lowerCamelCase_	,      lowerCamelCase_       	=		self.model_tester.prepare_config_and_inputs()
      for model_class in self.all_model_classes:
        lowerCamelCase_       	=		model_class(__SCREAMING_SNAKE_CASE   )

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






    def   UpperCamelCase			(			self	:				Tuple   )	->						str:
      lowerCamelCase_       	=		self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.check_encoder_decoder_model_standalone(*__SCREAMING_SNAKE_CASE   )






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






    def   UpperCamelCase			(			self	:				Tuple   )	->						Tuple:
      lowerCamelCase_       	=		inspect.signature(getattr(__SCREAMING_SNAKE_CASE			,		'forward'   )   )
      # The main input is the name of the argument after `self`
      lowerCamelCase_       	=		list(model_signature.parameters.keys()   )[1]
      self.assertEqual(AutoformerModel.main_input_name			,		__SCREAMING_SNAKE_CASE   )






    def   UpperCamelCase			(			self	:				Optional[Any]   )	->						str:
      lowerCamelCase_	,      lowerCamelCase_       	=		self.model_tester.prepare_config_and_inputs_for_common()

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

        lowerCamelCase_       	=		[
            '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(__SCREAMING_SNAKE_CASE   )]			,		__SCREAMING_SNAKE_CASE   )






    def   UpperCamelCase			(			self	:				List[str]   )	->						Optional[int]:
      lowerCamelCase_	,      lowerCamelCase_       	=		self.model_tester.prepare_config_and_inputs_for_common()
      lowerCamelCase_       	=		True

      lowerCamelCase_       	=		getattr(self.model_tester			,		'seq_length'			,		__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_       	=		getattr(self.model_tester			,		'decoder_seq_length'			,		__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_       	=		getattr(self.model_tester			,		'encoder_seq_length'			,		__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_       	=		getattr(self.model_tester			,		'd_model'			,		__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_       	=		getattr(self.model_tester			,		'num_attention_heads'			,		__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_       	=		d_model // num_attention_heads

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

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

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

        lowerCamelCase_       	=		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(__SCREAMING_SNAKE_CASE			,		__SCREAMING_SNAKE_CASE   )

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

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

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

      self.assertEqual(len(__SCREAMING_SNAKE_CASE   )			,		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	:				Any   )	->						List[Any]:
      super().test_retain_grad_hidden_states_attentions()







def 	lowerCamelCase__			(				_lowerCamelCase			:						List[str]="train-batch.pt"					) ->							Optional[Any]:
  lowerCamelCase_       	=		hf_hub_download(repo_id='hf-internal-testing/tourism-monthly-batch'						,							filename=_lowerCamelCase						,							repo_type='dataset'					)
  lowerCamelCase_       	=		torch.load(_lowerCamelCase						,							map_location=_lowerCamelCase					)
  return batch






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






    def   UpperCamelCase			(			self	:				Tuple   )	->						Optional[Any]:
      lowerCamelCase_       	=		AutoformerModel.from_pretrained('huggingface/autoformer-tourism-monthly'   ).to(__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_       	=		prepare_batch()

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

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

      lowerCamelCase_       	=		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=__SCREAMING_SNAKE_CASE   )
      self.assertTrue(torch.allclose(output[0, :3, :3]			,		__SCREAMING_SNAKE_CASE			,		atol=__SCREAMING_SNAKE_CASE   )   )






    def   UpperCamelCase			(			self	:				Tuple   )	->						Union[str, Any]:
      lowerCamelCase_       	=		AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly'   ).to(__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_       	=		prepare_batch('val-batch.pt'   )
      with torch.no_grad():
        lowerCamelCase_       	=		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
      lowerCamelCase_       	=		torch.Size((64, model.config.context_length, model.config.d_model)   )
      self.assertEqual(output.shape			,		__SCREAMING_SNAKE_CASE   )

      lowerCamelCase_       	=		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=__SCREAMING_SNAKE_CASE   )
      self.assertTrue(torch.allclose(output[0, :3, :3]			,		__SCREAMING_SNAKE_CASE			,		atol=__SCREAMING_SNAKE_CASE   )   )






    def   UpperCamelCase			(			self	:				List[str]   )	->						List[str]:
      lowerCamelCase_       	=		AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly'   ).to(__SCREAMING_SNAKE_CASE   )
      lowerCamelCase_       	=		prepare_batch('val-batch.pt'   )
      with torch.no_grad():
        lowerCamelCase_       	=		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']			,		)
      lowerCamelCase_       	=		torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length)   )
      self.assertEqual(outputs.sequences.shape			,		__SCREAMING_SNAKE_CASE   )

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












"""simple docstring"""





_SCREAMING_SNAKE_CASE   :							Optional[int]          =		8.314_4598







def 	lowerCamelCase__			(				_lowerCamelCase			:						float						,							_lowerCamelCase			:						float					) ->							float:
  if temperature < 0:
    raise Exception('Temperature cannot be less than 0 K'					)
  if molar_mass <= 0:
    raise Exception('Molar mass cannot be less than or equal to 0 kg/mol'					)
  else:
    return (3 * UNIVERSAL_GAS_CONSTANT * temperature / molar_mass) ** 0.5


if __name__ == "__main__":
  import doctest

  # run doctest
  doctest.testmod()

  # example
  _SCREAMING_SNAKE_CASE   :							List[Any]          =		300
  _SCREAMING_SNAKE_CASE   :							Dict          =		28
  _SCREAMING_SNAKE_CASE   :							List[str]          =		rms_speed_of_molecule(temperature, molar_mass)
  print(F'''Vrms of Nitrogen gas at 300 K is {vrms} m/s''')




