codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch autoformer model import inspect import tempfile import unittest from huggingfacehub import hfhubdownload from transformers import istorchavailable from transformers testingutils import isflaky requiretorch slow torchdevice from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin floatstensor idstensor from testpipelinemixin import pipelinetestermixin tolerance 1e4 if istorchavailable import torch from transformers import autoformerconfig autoformerforprediction autoformermodel from transformers models autoformer modelingautoformer import autoformerdecoder autoformerencoder requiretorch class autoformermodeltester def init self parent dmodel16 batchsize13 predictionlength7 contextlength14 labellength10 cardinality19 embeddingdimension5 numtimefeatures4 istrainingtrue hiddensize16 numhiddenlayers2 numattentionheads4 intermediatesize4 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 lagssequence1 2 3 4 5 movingaverage25 autocorrelationfactor5 self dmodel dmodel self parent parent self batchsize batchsize self predictionlength predictionlength self contextlength contextlength self cardinality cardinality self numtimefeatures numtimefeatures self lagssequence lagssequence self embeddingdimension embeddingdimension self istraining istraining self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self hiddenact hiddenact self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self encoderseqlength contextlength self decoderseqlength predictionlength labellength self labellength labellength self movingaverage movingaverage self autocorrelationfactor autocorrelationfactor def getconfigself return autoformerconfig dmodelself dmodel encoderlayersself numhiddenlayers decoderlayersself numhiddenlayers encoderattentionheadsself numattentionheads decoderattentionheadsself numattentionheads encoderffndimself intermediatesize decoderffndimself intermediatesize dropoutself hiddendropoutprob attentiondropoutself attentionprobsdropoutprob predictionlengthself predictionlength contextlengthself contextlength labellengthself labellength lagssequenceself lagssequence numtimefeaturesself numtimefeatures numstaticcategoricalfeatures1 cardinalityself cardinality embeddingdimensionself embeddingdimension movingaverageself movingaverage def prepareautoformerinputsdictself config pastlength config contextlength maxconfig lagssequence staticcategoricalfeatures idstensorself batchsize 1 config cardinality0 pasttimefeatures floatstensorself batchsize pastlength config numtimefeatures pastvalues floatstensorself batchsize pastlength pastobservedmask floatstensorself batchsize pastlength 0 5 decoder inputs futuretimefeatures floatstensorself batchsize config predictionlength config numtimefeatures futurevalues floatstensorself batchsize config predictionlength inputsdict pastvalues pastvalues staticcategoricalfeatures staticcategoricalfeatures pasttimefeatures pasttimefeatures pastobservedmask pastobservedmask futuretimefeatures futuretimefeatures futurevalues futurevalues return inputsdict def prepareconfigandinputsself config self getconfig inputsdict self prepareautoformerinputsdictconfig return config inputsdict def prepareconfigandinputsforcommonself config inputsdict self prepareconfigandinputs return config inputsdict def checkencoderdecodermodelstandaloneself config inputsdict model autoformermodelconfigconfig totorchdevice eval outputs modelinputsdict encoderlasthiddenstate outputs encoderlasthiddenstate lasthiddenstate outputs lasthiddenstate with tempfile temporarydirectory as tmpdirname encoder model getencoder encoder savepretrainedtmpdirname encoder autoformerencoder frompretrainedtmpdirname totorchdevice transformerinputs feature model createnetworkinputsinputsdict seasonalinput trendinput model decompositionlayertransformerinputs config contextlength encinput torch cat transformerinputs config contextlength feature config contextlength dim1 encoderlasthiddenstate2 encoderinputsembedsencinput0 self parent asserttrueencoderlasthiddenstate2 encoderlasthiddenstate abs max item 1e3 mean torch meantransformerinputs config contextlength dim1 unsqueeze1 repeat1 config predictionlength 1 zeros torch zeros transformerinputs shape0 config predictionlength transformerinputs shape2 deviceencinput device decinput torch cat torch catseasonalinput config labellength zeros dim1 feature config contextlength config labellength dim1 trendinit torch cat torch cattrendinput config labellength mean dim1 feature config contextlength config labellength dim1 with tempfile temporarydirectory as tmpdirname decoder model getdecoder decoder savepretrainedtmpdirname decoder autoformerdecoder frompretrainedtmpdirname totorchdevice lasthiddenstate2 decoder trendtrendinit inputsembedsdecinput encoderhiddenstatesencoderlasthiddenstate 0 self parent asserttruelasthiddenstate2 lasthiddenstate abs max item 1e3 requiretorch class autoformermodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses autoformermodel autoformerforprediction if istorchavailable else allgenerativemodelclasses autoformerforprediction if istorchavailable else pipelinemodelmapping featureextraction autoformermodel if istorchavailable else testpruning false testheadmasking false testmissingkeys false testtorchscript false testinputsembeds false testmodelcommonattributes false def setupself self modeltester autoformermodeltesterself self configtester configtesterself configclassautoformerconfig hastextmodalityfalse def testconfigself self configtester runcommontests def testsaveloadstrictself config inputsdict self modeltester prepareconfigandinputs for modelclass in self allmodelclasses model modelclassconfig with tempfile temporarydirectory as tmpdirname model savepretrainedtmpdirname model2 info modelclass frompretrainedtmpdirname outputloadinginfotrue self assertequalinfomissingkeys def testencoderdecodermodelstandaloneself configandinputs self modeltester prepareconfigandinputsforcommon self modeltester checkencoderdecodermodelstandaloneconfigandinputs unittest skipreasonmodel has no tokens embeddings def testresizetokensembeddingsself pass unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingself pass unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingusereentrantfalseself pass unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingusereentrantself pass input is staticcategoricalfeatures not inputids def testmodelmaininputnameself modelsignature inspect signaturegetattrautoformermodel forward the main input is the name of the argument after self observedmaininputname listmodelsignature parameters keys1 self assertequalautoformermodel maininputname observedmaininputname def testforwardsignatureself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig signature inspect signaturemodel forward signature parameters is an ordereddict so argnames order is deterministic argnames signature parameters keys expectedargnames pastvalues pasttimefeatures pastobservedmask staticcategoricalfeatures staticrealfeatures futurevalues futuretimefeatures if model class name in autoformerforprediction expectedargnames appendfutureobservedmask expectedargnames extend decoderattentionmask headmask decoderheadmask crossattnheadmask encoderoutputs pastkeyvalues outputhiddenstates outputattentions usecache returndict self assertlistequalargnames lenexpectedargnames expectedargnames def testattentionoutputsself config inputsdict self modeltester prepareconfigandinputsforcommon config returndict true seqlen getattrself modeltester seqlength none decoderseqlength getattrself modeltester decoderseqlength seqlen encoderseqlength getattrself modeltester encoderseqlength seqlen dmodel getattrself modeltester dmodel none numattentionheads getattrself modeltester numattentionheads none dim dmodel numattentionheads for modelclass in self allmodelclasses inputsdictoutputattentions true inputsdictoutputhiddenstates false config returndict true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs encoderattentions if config isencoderdecoder else outputs attentions self assertequallenattentions self modeltester numhiddenlayers check that outputattentions also work using config del inputsdictoutputattentions config outputattentions true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs encoderattentions self assertequallenattentions self modeltester numhiddenlayers self assertlistequal listattentions0 shape3 self modeltester numattentionheads encoderseqlength dim outlen lenoutputs correctoutlen 7 if lasthiddenstate in outputs correctoutlen 1 if trend in outputs correctoutlen 1 if pastkeyvalues in outputs correctoutlen 1 pastkeyvalues have been returned if loss in outputs correctoutlen 1 if params in outputs correctoutlen 1 self assertequaloutlen correctoutlen decoder attentions decoderattentions outputs decoderattentions self assertisinstancedecoderattentions list tuple self assertequallendecoderattentions self modeltester numhiddenlayers self assertlistequal listdecoderattentions0 shape3 self modeltester numattentionheads decoderseqlength dim cross attentions crossattentions outputs crossattentions self assertisinstancecrossattentions list tuple self assertequallencrossattentions self modeltester numhiddenlayers self assertlistequal listcrossattentions0 shape3 self modeltester numattentionheads decoderseqlength dim check attention is always last and order is fine inputsdictoutputattentions true inputsdictoutputhiddenstates true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass self assertequaloutlen 2 lenoutputs selfattentions outputs encoderattentions if config isencoderdecoder else outputs attentions self assertequallenselfattentions self modeltester numhiddenlayers self assertlistequal listselfattentions0 shape3 self modeltester numattentionheads encoderseqlength dim isflaky def testretaingradhiddenstatesattentionsself super testretaingradhiddenstatesattentions def preparebatchfilenametrainbatch pt file hfhubdownloadrepoidhfinternaltestingtourismmonthlybatch filenamefilename repotypedataset batch torch loadfile maplocationtorchdevice return batch requiretorch slow class autoformermodelintegrationtestsunittest testcase def testinferencenoheadself model autoformermodel frompretrainedhuggingfaceautoformertourismmonthly totorchdevice batch preparebatch with torch nograd output model pastvaluesbatchpastvalues pasttimefeaturesbatchpasttimefeatures pastobservedmaskbatchpastobservedmask staticcategoricalfeaturesbatchstaticcategoricalfeatures futurevaluesbatchfuturevalues futuretimefeaturesbatchfuturetimefeatures 0 expectedshape torch size 64 model config predictionlength model config labellength model config featuresize self assertequaloutput shape expectedshape expectedslice torch tensor 0 3593 1 3398 0 6330 0 2279 1 5396 0 1792 0 0450 1 3225 0 2335 devicetorchdevice self asserttruetorch allcloseoutput0 3 3 expectedslice atoltolerance def testinferenceheadself model autoformerforprediction frompretrainedhuggingfaceautoformertourismmonthly totorchdevice batch preparebatchvalbatch pt with torch nograd output model pastvaluesbatchpastvalues pasttimefeaturesbatchpasttimefeatures pastobservedmaskbatchpastobservedmask staticcategoricalfeaturesbatchstaticcategoricalfeatures encoderlasthiddenstate expectedshape torch size64 model config contextlength model config dmodel self assertequaloutput shape expectedshape expectedslice torch tensor 0 0734 0 9036 0 8358 4 7186 2 4113 1 9581 1 7953 2 3558 1 2970 devicetorchdevice self asserttruetorch allcloseoutput0 3 3 expectedslice atoltolerance def testseqtoseqgenerationself model autoformerforprediction frompretrainedhuggingfaceautoformertourismmonthly totorchdevice batch preparebatchvalbatch pt with torch nograd outputs model generate staticcategoricalfeaturesbatchstaticcategoricalfeatures pasttimefeaturesbatchpasttimefeatures pastvaluesbatchpastvalues futuretimefeaturesbatchfuturetimefeatures pastobservedmaskbatchpastobservedmask expectedshape torch size64 model config numparallelsamples model config predictionlength self assertequaloutputs sequences shape expectedshape expectedslice torch tensor3130 6763 4056 5293 7053 0786 devicetorchdevice meanprediction outputs sequences meandim1 self asserttruetorch allclosemeanprediction0 3 expectedslice rtol1e1 coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch autoformer model decoder inputs input is static_categorical_features not input_ids the main input is the name of the argument after self signature parameters is an ordereddict so arg_names order is deterministic check that output_attentions also work using config past_key_values have been returned decoder attentions cross attentions check attention is always last and order is fine
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


TOLERANCE = 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 AutoformerModelTester:
    def __init__(
        self,
        parent,
        d_model=16,
        batch_size=13,
        prediction_length=7,
        context_length=14,
        label_length=10,
        cardinality=19,
        embedding_dimension=5,
        num_time_features=4,
        is_training=True,
        hidden_size=16,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_size=4,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        lags_sequence=[1, 2, 3, 4, 5],
        moving_average=25,
        autocorrelation_factor=5,
    ):
        self.d_model = d_model
        self.parent = parent
        self.batch_size = batch_size
        self.prediction_length = prediction_length
        self.context_length = context_length
        self.cardinality = cardinality
        self.num_time_features = num_time_features
        self.lags_sequence = lags_sequence
        self.embedding_dimension = embedding_dimension
        self.is_training = is_training
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob

        self.encoder_seq_length = context_length
        self.decoder_seq_length = prediction_length + label_length
        self.label_length = label_length

        self.moving_average = moving_average
        self.autocorrelation_factor = autocorrelation_factor

    def get_config(self):
        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 prepare_autoformer_inputs_dict(self, config):
        _past_length = config.context_length + max(config.lags_sequence)

        static_categorical_features = ids_tensor([self.batch_size, 1], config.cardinality[0])
        past_time_features = floats_tensor([self.batch_size, _past_length, config.num_time_features])
        past_values = floats_tensor([self.batch_size, _past_length])
        past_observed_mask = floats_tensor([self.batch_size, _past_length]) > 0.5

        
        future_time_features = floats_tensor([self.batch_size, config.prediction_length, config.num_time_features])
        future_values = floats_tensor([self.batch_size, config.prediction_length])

        inputs_dict = {
            "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 prepare_config_and_inputs(self):
        config = self.get_config()
        inputs_dict = self.prepare_autoformer_inputs_dict(config)
        return config, inputs_dict

    def prepare_config_and_inputs_for_common(self):
        config, inputs_dict = self.prepare_config_and_inputs()
        return config, inputs_dict

    def check_encoder_decoder_model_standalone(self, config, inputs_dict):
        model = AutoformerModel(config=config).to(torch_device).eval()
        outputs = model(**inputs_dict)

        encoder_last_hidden_state = outputs.encoder_last_hidden_state
        last_hidden_state = outputs.last_hidden_state

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

        transformer_inputs, feature, _, _, _ = model.create_network_inputs(**inputs_dict)
        seasonal_input, trend_input = model.decomposition_layer(transformer_inputs[:, : config.context_length, ...])

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

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

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

        last_hidden_state_2 = decoder(
            trend=trend_init,
            inputs_embeds=dec_input,
            encoder_hidden_states=encoder_last_hidden_state,
        )[0]

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


@require_torch
class AutoformerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
    all_generative_model_classes = (AutoformerForPrediction,) if is_torch_available() else ()
    pipeline_model_mapping = {"feature-extraction": AutoformerModel} if is_torch_available() else {}
    test_pruning = False
    test_head_masking = False
    test_missing_keys = False
    test_torchscript = False
    test_inputs_embeds = False
    test_model_common_attributes = False

    def setUp(self):
        self.model_tester = AutoformerModelTester(self)
        self.config_tester = ConfigTester(self, config_class=AutoformerConfig, has_text_modality=False)

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_save_load_strict(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs()
        for model_class in self.all_model_classes:
            model = model_class(config)

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

    def test_encoder_decoder_model_standalone(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.check_encoder_decoder_model_standalone(*config_and_inputs)

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

    @unittest.skip(
        reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
    )
    def test_training_gradient_checkpointing(self):
        pass

    @unittest.skip(
        reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
    )
    def test_training_gradient_checkpointing_use_reentrant_false(self):
        pass

    @unittest.skip(
        reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
    )
    def test_training_gradient_checkpointing_use_reentrant(self):
        pass

    
    def test_model_main_input_name(self):
        model_signature = inspect.signature(getattr(AutoformerModel, "forward"))
        
        observed_main_input_name = list(model_signature.parameters.keys())[1]
        self.assertEqual(AutoformerModel.main_input_name, observed_main_input_name)

    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.forward)
            
            arg_names = [*signature.parameters.keys()]

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

    def test_attention_outputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.return_dict = True

        seq_len = getattr(self.model_tester, "seq_length", None)
        decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", seq_len)
        encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", seq_len)
        d_model = getattr(self.model_tester, "d_model", None)
        num_attention_heads = getattr(self.model_tester, "num_attention_heads", None)
        dim = d_model // num_attention_heads

        for model_class in self.all_model_classes:
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = False
            config.return_dict = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

            
            del inputs_dict["output_attentions"]
            config.output_attentions = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.encoder_attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

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

            correct_outlen = 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  

            if "loss" in outputs:
                correct_outlen += 1

            if "params" in outputs:
                correct_outlen += 1

            self.assertEqual(out_len, correct_outlen)

            
            decoder_attentions = outputs.decoder_attentions
            self.assertIsInstance(decoder_attentions, (list, tuple))
            self.assertEqual(len(decoder_attentions), 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 = outputs.cross_attentions
            self.assertIsInstance(cross_attentions, (list, tuple))
            self.assertEqual(len(cross_attentions), self.model_tester.num_hidden_layers)
            self.assertListEqual(
                list(cross_attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads, decoder_seq_length, dim],
            )

        
        inputs_dict["output_attentions"] = True
        inputs_dict["output_hidden_states"] = True
        model = model_class(config)
        model.to(torch_device)
        model.eval()
        with torch.no_grad():
            outputs = model(**self._prepare_for_class(inputs_dict, model_class))

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

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

        self.assertEqual(len(self_attentions), 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 test_retain_grad_hidden_states_attentions(self):
        super().test_retain_grad_hidden_states_attentions()


def prepare_batch(filename="train-batch.pt"):
    file = hf_hub_download(repo_id="hf-internal-testing/tourism-monthly-batch", filename=filename, repo_type="dataset")
    batch = torch.load(file, map_location=torch_device)
    return batch


@require_torch
@slow
class AutoformerModelIntegrationTests(unittest.TestCase):
    def test_inference_no_head(self):
        model = AutoformerModel.from_pretrained("huggingface/autoformer-tourism-monthly").to(torch_device)
        batch = prepare_batch()

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

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

        expected_slice = torch.tensor(
            [[0.3593, -1.3398, 0.6330], [0.2279, 1.5396, -0.1792], [0.0450, 1.3225, -0.2335]], device=torch_device
        )
        self.assertTrue(torch.allclose(output[0, :3, :3], expected_slice, atol=TOLERANCE))

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

        expected_slice = torch.tensor(
            [[-0.0734, -0.9036, 0.8358], [4.7186, 2.4113, 1.9581], [1.7953, 2.3558, 1.2970]], device=torch_device
        )
        self.assertTrue(torch.allclose(output[0, :3, :3], expected_slice, atol=TOLERANCE))

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

        expected_slice = torch.tensor([3130.6763, 4056.5293, 7053.0786], device=torch_device)
        mean_prediction = outputs.sequences.mean(dim=1)
        self.assertTrue(torch.allclose(mean_prediction[0, -3:], expected_slice, rtol=1e-1))