

'''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


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

     SCREAMING_SNAKE_CASE:List[Any]  					=    context_length
     SCREAMING_SNAKE_CASE:Optional[Any]  					=    prediction_length + label_length
     SCREAMING_SNAKE_CASE:Any  					=    label_length

     SCREAMING_SNAKE_CASE:Tuple  					=    moving_average
     SCREAMING_SNAKE_CASE:int  					=    autocorrelation_factor
   def 			__UpperCamelCase  (       self	:					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	:					Optional[int]   ,SCREAMING_SNAKE_CASE__	:					Optional[Any]							):
     SCREAMING_SNAKE_CASE:Optional[int]  					=    config.context_length + max(config.lags_sequence							)

     SCREAMING_SNAKE_CASE:Optional[Any]  					=    ids_tensor([self.batch_size, 1]   ,config.cardinality[0]							)
     SCREAMING_SNAKE_CASE:List[str]  					=    floats_tensor([self.batch_size, _past_length, config.num_time_features]							)
     SCREAMING_SNAKE_CASE:Dict  					=    floats_tensor([self.batch_size, _past_length]							)
     SCREAMING_SNAKE_CASE:Any  					=    floats_tensor([self.batch_size, _past_length]							) > 0.5

     # decoder inputs
     SCREAMING_SNAKE_CASE:Union[str, Any]  					=    floats_tensor([self.batch_size, config.prediction_length, config.num_time_features]							)
     SCREAMING_SNAKE_CASE:str  					=    floats_tensor([self.batch_size, config.prediction_length]							)

     SCREAMING_SNAKE_CASE:Optional[int]  					=    {
         "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	:					List[Any]							):
     SCREAMING_SNAKE_CASE:Optional[int]  					=    self.get_config()
     SCREAMING_SNAKE_CASE:List[str]  					=    self.prepare_autoformer_inputs_dict(SCREAMING_SNAKE_CASE__							)
     return config, inputs_dict
   def 			__UpperCamelCase  (       self	:					Optional[int]							):
     SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE:Optional[int]  					=    self.prepare_config_and_inputs()
     return config, inputs_dict







   def 			__UpperCamelCase  (       self	:					Optional[Any]   ,SCREAMING_SNAKE_CASE__	:					Any   ,SCREAMING_SNAKE_CASE__	:					str							):
     SCREAMING_SNAKE_CASE:str  					=    AutoformerModel(config=SCREAMING_SNAKE_CASE__							).to(SCREAMING_SNAKE_CASE__							).eval()
     SCREAMING_SNAKE_CASE:Tuple  					=    model(**SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Union[str, Any]  					=    outputs.encoder_last_hidden_state
     SCREAMING_SNAKE_CASE:Optional[int]  					=    outputs.last_hidden_state

     with tempfile.TemporaryDirectory() as tmpdirname:
       SCREAMING_SNAKE_CASE:Optional[Any]  					=    model.get_encoder()
       encoder.save_pretrained(SCREAMING_SNAKE_CASE__							)
       SCREAMING_SNAKE_CASE:str  					=    AutoformerEncoder.from_pretrained(SCREAMING_SNAKE_CASE__							).to(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE:Tuple  					=    model.create_network_inputs(**SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE:Union[str, Any]  					=    model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]							)

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

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

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

     SCREAMING_SNAKE_CASE:Optional[int]  					=    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 							_snake_case       (				_a		,     _a		,     unittest.TestCase  ):
   _A				:   Dict			      =    (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
   _A				:   Optional[Any]			      =    (AutoformerForPrediction,) if is_torch_available() else ()
   _A				:   Any			      =    {'feature-extraction': AutoformerModel} if is_torch_available() else {}
   _A				:   Dict			      =    False
   _A				:   Any			      =    False
   _A				:   Optional[Any]			      =    False
   _A				:   Dict			      =    False
   _A				:   int			      =    False
   _A				:   int			      =    False
   def 			__UpperCamelCase  (       self	:					Any							):
     SCREAMING_SNAKE_CASE:Union[str, Any]  					=    AutoformerModelTester(self							)
     SCREAMING_SNAKE_CASE:Any  					=    ConfigTester(self   ,config_class=SCREAMING_SNAKE_CASE__   ,has_text_modality=SCREAMING_SNAKE_CASE__							)
   def 			__UpperCamelCase  (       self	:					str							):
     self.config_tester.run_common_tests()
   def 			__UpperCamelCase  (       self	:					Any							):
     SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE:str  					=    self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE:Union[str, Any]  					=    model_class(SCREAMING_SNAKE_CASE__							)

       with tempfile.TemporaryDirectory() as tmpdirname:
         model.save_pretrained(SCREAMING_SNAKE_CASE__							)
         SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE:List[str]  					=    model_class.from_pretrained(SCREAMING_SNAKE_CASE__   ,output_loading_info=SCREAMING_SNAKE_CASE__							)
       self.assertEqual(info["missing_keys"]   ,[]							)
   def 			__UpperCamelCase  (       self	:					Any							):
     SCREAMING_SNAKE_CASE:List[str]  					=    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	:					Dict							):
     pass
   def 			__UpperCamelCase  (       self	:					str							):
     SCREAMING_SNAKE_CASE:Optional[Any]  					=    inspect.signature(getattr(SCREAMING_SNAKE_CASE__   ,"forward"							)							)
     # The main input is the name of the argument after `self`
     SCREAMING_SNAKE_CASE:List[str]  					=    list(model_signature.parameters.keys()							)[1]
     self.assertEqual(AutoformerModel.main_input_name   ,SCREAMING_SNAKE_CASE__							)
   def 			__UpperCamelCase  (       self	:					List[str]							):
     SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE:Optional[int]  					=    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE:Tuple  					=    model_class(SCREAMING_SNAKE_CASE__							)
       SCREAMING_SNAKE_CASE:List[Any]  					=    inspect.signature(model.forward							)
       # signature.parameters is an OrderedDict => so arg_names order is deterministic
       SCREAMING_SNAKE_CASE:List[str]  					=    [*signature.parameters.keys()]

       SCREAMING_SNAKE_CASE:Dict  					=    [
           "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[Any]							):
     SCREAMING_SNAKE_CASE					,			SCREAMING_SNAKE_CASE:Tuple  					=    self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE:List[str]  					=    True

     SCREAMING_SNAKE_CASE:Optional[Any]  					=    getattr(self.model_tester   ,"seq_length"   ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Optional[int]  					=    getattr(self.model_tester   ,"decoder_seq_length"   ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Dict  					=    getattr(self.model_tester   ,"encoder_seq_length"   ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:int  					=    getattr(self.model_tester   ,"d_model"   ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:List[Any]  					=    getattr(self.model_tester   ,"num_attention_heads"   ,SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Optional[Any]  					=    d_model // num_attention_heads

     for model_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE:Any  					=    True
       SCREAMING_SNAKE_CASE:Dict  					=    False
       SCREAMING_SNAKE_CASE:Any  					=    True
       SCREAMING_SNAKE_CASE:Any  					=    model_class(SCREAMING_SNAKE_CASE__							)
       model.to(SCREAMING_SNAKE_CASE__							)
       model.eval()
       with torch.no_grad():
         SCREAMING_SNAKE_CASE:Optional[int]  					=    model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__   ,SCREAMING_SNAKE_CASE__							)							)
       SCREAMING_SNAKE_CASE:Dict  					=    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"]
       SCREAMING_SNAKE_CASE:int  					=    True
       SCREAMING_SNAKE_CASE:Tuple  					=    model_class(SCREAMING_SNAKE_CASE__							)
       model.to(SCREAMING_SNAKE_CASE__							)
       model.eval()
       with torch.no_grad():
         SCREAMING_SNAKE_CASE:Any  					=    model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__   ,SCREAMING_SNAKE_CASE__							)							)
       SCREAMING_SNAKE_CASE:List[str]  					=    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]   ,)
       SCREAMING_SNAKE_CASE:Optional[int]  					=    len(SCREAMING_SNAKE_CASE__							)

       SCREAMING_SNAKE_CASE:Optional[int]  					=    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
       SCREAMING_SNAKE_CASE:str  					=    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
       SCREAMING_SNAKE_CASE:Dict  					=    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
     SCREAMING_SNAKE_CASE:Dict  					=    True
     SCREAMING_SNAKE_CASE:Dict  					=    True
     SCREAMING_SNAKE_CASE:List[str]  					=    model_class(SCREAMING_SNAKE_CASE__							)
     model.to(SCREAMING_SNAKE_CASE__							)
     model.eval()
     with torch.no_grad():
       SCREAMING_SNAKE_CASE:Union[str, Any]  					=    model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__   ,SCREAMING_SNAKE_CASE__							)							)

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

     SCREAMING_SNAKE_CASE:Any  					=    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	:					Union[str, Any]							):
     super().test_retain_grad_hidden_states_attentions()





def       A_     (       snake_case="train-batch.pt"       ):
  SCREAMING_SNAKE_CASE:Dict  					=    hf_hub_download(repo_id="hf-internal-testing/tourism-monthly-batch"					,				filename=__A					,				repo_type="dataset"       )
  SCREAMING_SNAKE_CASE:Optional[Any]  					=    torch.load(__A					,				map_location=__A       )
  return batch





@require_torch
@slow
class 							_snake_case       (				unittest.TestCase  ):
   def 			__UpperCamelCase  (       self	:					int							):
     SCREAMING_SNAKE_CASE:Optional[Any]  					=    AutoformerModel.from_pretrained("huggingface/autoformer-tourism-monthly"							).to(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:str  					=    prepare_batch()

     with torch.no_grad():
       SCREAMING_SNAKE_CASE:List[str]  					=    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]

     SCREAMING_SNAKE_CASE:List[str]  					=    torch.Size(
         (64, model.config.prediction_length + model.config.label_length, model.config.feature_size)							)
     self.assertEqual(output.shape   ,SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Dict  					=    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	:					int							):
     SCREAMING_SNAKE_CASE:int  					=    AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly"							).to(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Tuple  					=    prepare_batch("val-batch.pt"							)
     with torch.no_grad():
       SCREAMING_SNAKE_CASE:Any  					=    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
     SCREAMING_SNAKE_CASE:Tuple  					=    torch.Size((64, model.config.context_length, model.config.d_model)							)
     self.assertEqual(output.shape   ,SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Union[str, Any]  					=    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]							):
     SCREAMING_SNAKE_CASE:Union[str, Any]  					=    AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly"							).to(SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:str  					=    prepare_batch("val-batch.pt"							)
     with torch.no_grad():
       SCREAMING_SNAKE_CASE:Tuple  					=    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"]   ,)
     SCREAMING_SNAKE_CASE:Dict  					=    torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length)							)
     self.assertEqual(outputs.sequences.shape   ,SCREAMING_SNAKE_CASE__							)

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











'''simple docstring'''





from collections.abc import Generator
from math import sin




def lowerCAmelCase  (__A):



     """simple docstring"""



     if len(__A) != 32:
          raise ValueError('''Input must be of length 32''')

     _a      =  b''''''
     for i in [3, 2, 1, 0]:
          little_endian += string_aa[8 * i : 8 * i + 8]
     return little_endian




def lowerCAmelCase  (__A):



     """simple docstring"""



     if i < 0:
          raise ValueError('''Input must be non-negative''')

     _a      =  format(__A  ,	'''08x''')[-8:]
     _a      =  b''''''
     for i in [3, 2, 1, 0]:
          little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode('''utf-8''')
     return little_endian_hex




def lowerCAmelCase  (__A):



     """simple docstring"""



     _a      =  b''''''
     for char in message:
          bit_string += format(__A  ,	'''08b''').encode('''utf-8''')
     _a      =  format(len(__A)  ,	'''064b''').encode('''utf-8''')

     # Pad bit_string to a multiple of 512 chars
     bit_string += b"1"
     while len(__A) % 512 != 448:
          bit_string += b"0"
     bit_string += to_little_endian(start_len[32:]) + to_little_endian(start_len[:32])

     return bit_string




def lowerCAmelCase  (__A):



     """simple docstring"""



     if len(__A) % 512 != 0:
          raise ValueError('''Input must have length that\'s a multiple of 512''')

     for pos in range(0  ,	len(__A)  ,	512):
          _a      =  bit_string[pos : pos + 512]
          _a      =  []
          for i in range(0  ,	512  ,	32):
               block_words.append(int(to_little_endian(block[i : i + 32])  ,	2))
          yield block_words




def lowerCAmelCase  (__A):



     """simple docstring"""



     if i < 0:
          raise ValueError('''Input must be non-negative''')

     _a      =  format(__A  ,	'''032b''')
     _a      =  ''''''
     for c in i_str:
          new_str += "1" if c == "0" else "0"
     return int(__A  ,	2)




def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



     return (a + b) % 2**32




def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



     if i < 0:
          raise ValueError('''Input must be non-negative''')
     if shift < 0:
          raise ValueError('''Shift must be non-negative''')
     return ((i << shift) ^ (i >> (32 - shift))) % 2**32




def lowerCAmelCase  (__A):



     """simple docstring"""



     _a      =  preprocess(__A)

     _a      =  [int(2**32 * abs(sin(i + 1))) for i in range(64)]

     # Starting states
     _a      =  0x67_452_301
     _a      =  0xEF_CDA_B89
     _a      =  0x98_BAD_CFE
     _a      =  0x10_325_476

     _a      =  [
         7,
         12,
         17,
         22,
         7,
         12,
         17,
         22,
         7,
         12,
         17,
         22,
         7,
         12,
         17,
         22,
         5,
         9,
         14,
         20,
         5,
         9,
         14,
         20,
         5,
         9,
         14,
         20,
         5,
         9,
         14,
         20,
         4,
         11,
         16,
         23,
         4,
         11,
         16,
         23,
         4,
         11,
         16,
         23,
         4,
         11,
         16,
         23,
         6,
         10,
         15,
         21,
         6,
         10,
         15,
         21,
         6,
         10,
         15,
         21,
         6,
         10,
         15,
         21,
     ]

     # Process bit string in chunks, each with 16 32-char words
     for block_words in get_block_words(__A):
          _a      =  aa
          _a      =  ba
          _a      =  ca
          _a      =  da

          # Hash current chunk
          for i in range(64):
               if i <= 15:
                    # f = (b & c) | (not_32(b) & d)     # Alternate definition for f
                    _a      =  d ^ (b & (c ^ d))
                    _a      =  i
               elif i <= 31:
                    # f = (d & b) | (not_32(d) & c)     # Alternate definition for f
                    _a      =  c ^ (d & (b ^ c))
                    _a      =  (5 * i + 1) % 16
               elif i <= 47:
                    _a      =  b ^ c ^ d
                    _a      =  (3 * i + 5) % 16
               else:
                    _a      =  c ^ (b | not_aa(__A))
                    _a      =  (7 * i) % 16
               _a      =  (f + a + added_consts[i] + block_words[g]) % 2**32
               _a      =  d
               _a      =  c
               _a      =  b
               _a      =  sum_aa(__A  ,	left_rotate_aa(__A  ,	shift_amounts[i]))

          # Add hashed chunk to running total
          _a      =  sum_aa(__A  ,	__A)
          _a      =  sum_aa(__A  ,	__A)
          _a      =  sum_aa(__A  ,	__A)
          _a      =  sum_aa(__A  ,	__A)

     _a      =  reformat_hex(__A) + reformat_hex(__A) + reformat_hex(__A) + reformat_hex(__A)
     return digest


if __name__ == "__main__":
  import doctest

  doctest.testmod()



