





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 		SCREAMING_SNAKE_CASE__				:


      def __init__(self       :      Union[str, Any]				,			a__       :      Union[str, Any]				,			a__       :      Optional[Any]=16				,			a__       :      str=13				,			a__       :      str=7				,			a__       :      Optional[int]=14				,			a__       :      Tuple=10				,			a__       :      Dict=19				,			a__       :      Optional[Any]=5				,			a__       :      Any=4				,			a__       :      int=True				,			a__       :      List[str]=16				,			a__       :      Tuple=2				,			a__       :      Any=4				,			a__       :      Dict=4				,			a__       :      Optional[Any]="gelu"				,			a__       :      Tuple=0.1				,			a__       :      Any=0.1				,			a__       :      Optional[int]=[1, 2, 3, 4, 5]				,			a__       :      str=25				,			a__       :      Optional[Any]=5				,			):





         """simple docstring"""



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

         __snake_case							       =   context_length
         __snake_case							       =   prediction_length + label_length
         __snake_case							       =   label_length

         __snake_case							       =   moving_average
         __snake_case							       =   autocorrelation_factor


      def    a (self       :      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    a (self       :      Any				,			a__       :      Union[str, Any]      ):





         """simple docstring"""



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

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

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

         __snake_case							       =   {
             '''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    a (self       :      Tuple      ):





         """simple docstring"""



         __snake_case							       =   self.get_config()
         __snake_case							       =   self.prepare_autoformer_inputs_dict(a__      )
         return config, inputs_dict


      def    a (self       :      Optional[int]      ):





         """simple docstring"""



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


      def    a (self       :      Dict				,			a__       :      int				,			a__       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   AutoformerModel(config=a__      ).to(a__      ).eval()
         __snake_case							       =   model(**a__      )

         __snake_case							       =   outputs.encoder_last_hidden_state
         __snake_case							       =   outputs.last_hidden_state

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

         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   model.create_network_inputs(**a__      )
         __snake_case       ,	__snake_case							       =   model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]      )

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

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

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

         __snake_case							       =   decoder(
             trend=a__				,			inputs_embeds=a__				,			encoder_hidden_states=a__				,			)[0]

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






@require_torch
class 		SCREAMING_SNAKE_CASE__				(       _UpperCAmelCase						,      _UpperCAmelCase						,      unittest.TestCase  ):
      A_						:				Optional[Any]							=    (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
      A_						:				str							=    (AutoformerForPrediction,) if is_torch_available() else ()
      A_						:				Optional[int]							=    {'feature-extraction': AutoformerModel} if is_torch_available() else {}
      A_						:				List[Any]							=    False
      A_						:				List[str]							=    False
      A_						:				Any							=    False
      A_						:				Union[str, Any]							=    False
      A_						:				Dict							=    False
      A_						:				Dict							=    False


      def    a (self       :      Optional[Any]      ):





         """simple docstring"""



         __snake_case							       =   AutoformerModelTester(self      )
         __snake_case							       =   ConfigTester(self				,			config_class=a__				,			has_text_modality=a__      )


      def    a (self       :      Optional[int]      ):





         """simple docstring"""



         self.config_tester.run_common_tests()


      def    a (self       :      int      ):





         """simple docstring"""



         __snake_case       ,	__snake_case							       =   self.model_tester.prepare_config_and_inputs()
         for model_class in self.all_model_classes:
            __snake_case							       =   model_class(a__      )

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


      def    a (self       :      Optional[int]      ):





         """simple docstring"""



         __snake_case							       =   self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.check_encoder_decoder_model_standalone(*a__      )


      @unittest.skip(reason='''Model has no tokens embeddings'''      )
      def    a (self       :      str      ):





         """simple docstring"""



         pass


      def    a (self       :      Dict      ):





         """simple docstring"""



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


      def    a (self       :      Any      ):





         """simple docstring"""



         __snake_case       ,	__snake_case							       =   self.model_tester.prepare_config_and_inputs_for_common()

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

            __snake_case							       =   [
                '''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(a__      )]				,			a__      )


      def    a (self       :      str      ):





         """simple docstring"""



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

         __snake_case							       =   getattr(self.model_tester				,			'''seq_length'''				,			a__      )
         __snake_case							       =   getattr(self.model_tester				,			'''decoder_seq_length'''				,			a__      )
         __snake_case							       =   getattr(self.model_tester				,			'''encoder_seq_length'''				,			a__      )
         __snake_case							       =   getattr(self.model_tester				,			'''d_model'''				,			a__      )
         __snake_case							       =   getattr(self.model_tester				,			'''num_attention_heads'''				,			a__      )
         __snake_case							       =   d_model // num_attention_heads

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

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

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

            __snake_case							       =   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(a__				,			a__      )

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

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

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

         self.assertEqual(len(a__      )				,			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    a (self       :      str      ):





         """simple docstring"""



         super().test_retain_grad_hidden_states_attentions()






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







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


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   AutoformerModel.from_pretrained('''huggingface/autoformer-tourism-monthly'''      ).to(a__      )
         __snake_case							       =   prepare_batch()

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

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

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


      def    a (self       :      List[Any]      ):





         """simple docstring"""



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

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


      def    a (self       :      List[Any]      ):





         """simple docstring"""



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

         __snake_case							       =   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=a__      )
         __snake_case							       =   outputs.sequences.mean(dim=1      )
         self.assertTrue(torch.allclose(mean_prediction[0, -3:]				,			a__				,			rtol=1E-1      )      )








def lowerCamelCase__    (							snake_case_	:  int       )      ->				int:
   if not isinstance(snake_case_       ,      snake_case_       ) or number < 0:
      raise ValueError('''Input must be a non-negative integer'''       )

   __snake_case							       =   0
   while number:
      # This way we arrive at next set bit (next 1) instead of looping
      # through each bit and checking for 1s hence the
      # loop won't run 32 times it will only run the number of `1` times
      number &= number - 1
      count += 1
   return count


if __name__ == "__main__":
    import doctest

    doctest.testmod()

