codingutf8 2022 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 timeseriestransformer model import inspect import tempfile import unittest from huggingfacehub import hfhubdownload from parameterized import parameterized 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 timeseriestransformerconfig timeseriestransformerforprediction timeseriestransformermodel from transformers models timeseriestransformer modelingtimeseriestransformer import timeseriestransformerdecoder timeseriestransformerencoder requiretorch class timeseriestransformermodeltester def init self parent batchsize13 predictionlength7 contextlength14 cardinality19 embeddingdimension5 numtimefeatures4 istrainingtrue hiddensize64 numhiddenlayers2 numattentionheads4 intermediatesize4 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 lagssequence1 2 3 4 5 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 def getconfigself return timeseriestransformerconfig encoderlayersself numhiddenlayers decoderlayersself numhiddenlayers encoderattentionheadsself numattentionheads decoderattentionheadsself numattentionheads encoderffndimself intermediatesize decoderffndimself intermediatesize dropoutself hiddendropoutprob attentiondropoutself attentionprobsdropoutprob predictionlengthself predictionlength contextlengthself contextlength lagssequenceself lagssequence numtimefeaturesself numtimefeatures numstaticrealfeatures1 numstaticcategoricalfeatures1 cardinalityself cardinality embeddingdimensionself embeddingdimension def preparetimeseriestransformerinputsdictself config pastlength config contextlength maxconfig lagssequence staticcategoricalfeatures idstensorself batchsize 1 config cardinality0 staticrealfeatures floatstensorself batchsize 1 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 staticrealfeatures staticrealfeatures pasttimefeatures pasttimefeatures pastobservedmask pastobservedmask futuretimefeatures futuretimefeatures futurevalues futurevalues return inputsdict def prepareconfigandinputsself config self getconfig inputsdict self preparetimeseriestransformerinputsdictconfig return config inputsdict def prepareconfigandinputsforcommonself config inputsdict self prepareconfigandinputs return config inputsdict def checkencoderdecodermodelstandaloneself config inputsdict model timeseriestransformermodelconfigconfig totorchdevice eval outputs modelinputsdict encoderlasthiddenstate outputs encoderlasthiddenstate lasthiddenstate outputs lasthiddenstate with tempfile temporarydirectory as tmpdirname encoder model getencoder encoder savepretrainedtmpdirname encoder timeseriestransformerencoder frompretrainedtmpdirname totorchdevice transformerinputs model createnetworkinputsinputsdict encinput transformerinputs config contextlength decinput transformerinputs config contextlength encoderlasthiddenstate2 encoderinputsembedsencinput0 self parent asserttrueencoderlasthiddenstate2 encoderlasthiddenstate abs max item 1e3 with tempfile temporarydirectory as tmpdirname decoder model getdecoder decoder savepretrainedtmpdirname decoder timeseriestransformerdecoder frompretrainedtmpdirname totorchdevice lasthiddenstate2 decoder inputsembedsdecinput encoderhiddenstatesencoderlasthiddenstate 0 self parent asserttruelasthiddenstate2 lasthiddenstate abs max item 1e3 requiretorch class timeseriestransformermodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses timeseriestransformermodel timeseriestransformerforprediction if istorchavailable else allgenerativemodelclasses timeseriestransformerforprediction if istorchavailable else pipelinemodelmapping featureextraction timeseriestransformermodel if istorchavailable else isencoderdecoder true testpruning false testheadmasking false testmissingkeys false testtorchscript false testinputsembeds false testmodelcommonattributes false def setupself self modeltester timeseriestransformermodeltesterself self configtester configtester self configclasstimeseriestransformerconfig hastextmodalityfalse predictionlengthself modeltester predictionlength def testconfigself self configtester runcommontests def testsaveloadstrictself config 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 ignore since we have no tokens embeddings def testresizetokensembeddingsself pass input is staticcategoricalfeatures not inputids def testmodelmaininputnameself modelsignature inspect signaturegetattrtimeseriestransformermodel forward the main input is the name of the argument after self observedmaininputname listmodelsignature parameters keys1 self assertequaltimeseriestransformermodel 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 expectedargnames extend futureobservedmask decoderattentionmask headmask decoderheadmask crossattnheadmask encoderoutputs pastkeyvalues outputhiddenstates outputattentions usecache returndict if futureobservedmask in argnames else 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 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 encoderseqlength outlen lenoutputs correctoutlen 7 if lasthiddenstate 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 decoderseqlength cross attentions crossattentions outputs crossattentions self assertisinstancecrossattentions list tuple self assertequallencrossattentions self modeltester numhiddenlayers self assertlistequal listcrossattentions0 shape3 self modeltester numattentionheads decoderseqlength encoderseqlength 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 encoderseqlength 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 testtraininggradientcheckpointingusereentrantself pass unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingusereentrantfalseself pass parameterized expand 1 5 1 1 5 1 10 15 1 5 3 6 9 10 2 5 1 2 7 2 5 2 3 4 6 4 5 1 5 9 11 4 5 7 8 13 14 def testcreatenetworkinputsself predictionlength contextlength lagssequence historylength maxlagssequence contextlength config timeseriestransformerconfig predictionlengthpredictionlength contextlengthcontextlength lagssequencelagssequence scalingfalse numparallelsamples10 numstaticcategoricalfeatures1 cardinality1 embeddingdimension2 numstaticrealfeatures1 model timeseriestransformermodelconfig batch staticcategoricalfeatures torch tensor0 dtypetorch int64 staticrealfeatures torch tensor0 0 dtypetorch float32 pasttimefeatures torch arangehistorylength dtypetorch float32 view1 historylength 1 pastvalues torch arangehistorylength dtypetorch float32 view1 historylength pastobservedmask torch arangehistorylength dtypetorch float32 view1 historylength test with no futuretarget only one step prediction batchfuturetimefeatures torch arangehistorylength historylength 1 dtypetorch float32 view 1 1 1 transformerinputs loc scale model createnetworkinputsbatch self asserttruescale 1 0 all assert loc 0 0 all ref torch arangemaxlagssequence historylength dtypetorch float32 for idx lag in enumeratelagssequence assert torch iscloseref lag transformerinputs0 idx all test with all future data batchfuturetimefeatures torch arange historylength historylength predictionlength dtypetorch float32 view1 predictionlength 1 batchfuturevalues torch arange historylength historylength predictionlength dtypetorch float32 view1 predictionlength transformerinputs loc scale model createnetworkinputsbatch assert scale 1 0 all assert loc 0 0 all ref torch arangemaxlagssequence historylength predictionlength dtypetorch float32 for idx lag in enumeratelagssequence assert torch iscloseref lag transformerinputs0 idx all test for generation batch popfuturevalues transformerinputs loc scale model createnetworkinputsbatch laggedsequence model getlaggedsubsequences sequencebatchpastvalues subsequenceslength1 shift1 assert that the last element of the lagged sequence is the one after the encoders input assert transformerinputs0 01 1 laggedsequence0 01 futurevalues torch arangehistorylength historylength predictionlength dtypetorch float32 view 1 predictionlength assert that the first element of the futurevalues is offset by lag after the decoders input assert laggedsequence0 01 lagssequence0 futurevalues0 0 isflaky def testretaingradhiddenstatesattentionsself super testretaingradhiddenstatesattentions def preparebatchfilenametrainbatch pt file hfhubdownloadrepoidhfinternaltestingtourismmonthlybatch filenamefilename repotypedataset batch torch loadfile maplocationtorchdevice return batch requiretorch slow class timeseriestransformermodelintegrationtestsunittest testcase def testinferencenoheadself model timeseriestransformermodel frompretrainedhuggingfacetimeseriestransformertourismmonthly to torchdevice batch preparebatch with torch nograd output model pastvaluesbatchpastvalues pasttimefeaturesbatchpasttimefeatures pastobservedmaskbatchpastobservedmask staticcategoricalfeaturesbatchstaticcategoricalfeatures staticrealfeaturesbatchstaticrealfeatures futurevaluesbatchfuturevalues futuretimefeaturesbatchfuturetimefeatures lasthiddenstate expectedshape torch size64 model config contextlength model config dmodel self assertequaloutput shape expectedshape expectedslice torch tensor 0 8196 1 5131 1 4620 1 1268 1 3238 1 5997 1 5098 1 0715 1 7359 devicetorchdevice self asserttruetorch allcloseoutput0 3 3 expectedslice atoltolerance def testinferenceheadself model timeseriestransformerforprediction frompretrained huggingfacetimeseriestransformertourismmonthly totorchdevice batch preparebatchvalbatch pt with torch nograd output model pastvaluesbatchpastvalues pasttimefeaturesbatchpasttimefeatures pastobservedmaskbatchpastobservedmask staticcategoricalfeaturesbatchstaticcategoricalfeatures staticrealfeaturesbatchstaticrealfeatures futuretimefeaturesbatchfuturetimefeatures encoderlasthiddenstate expectedshape torch size64 model config contextlength model config dmodel self assertequaloutput shape expectedshape expectedslice torch tensor 1 2957 1 0280 0 6045 0 7017 0 8193 0 3717 1 0449 0 8149 0 1405 devicetorchdevice self asserttruetorch allcloseoutput0 3 3 expectedslice atoltolerance def testseqtoseqgenerationself model timeseriestransformerforprediction frompretrained huggingfacetimeseriestransformertourismmonthly totorchdevice batch preparebatchvalbatch pt with torch nograd outputs model generate staticcategoricalfeaturesbatchstaticcategoricalfeatures staticrealfeaturesbatchstaticrealfeatures pasttimefeaturesbatchpasttimefeatures pastvaluesbatchpastvalues futuretimefeaturesbatchfuturetimefeatures pastobservedmaskbatchpastobservedmask expectedshape torch size64 model config numparallelsamples model config predictionlength self assertequaloutputs sequences shape expectedshape expectedslice torch tensor2825 2749 3584 9207 6763 9951 devicetorchdevice meanprediction outputs sequences meandim1 self asserttruetorch allclosemeanprediction0 3 expectedslice rtol1e1 coding utf 8 2022 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 timeseriestransformer model decoder inputs ignore since we have no tokens embeddings 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 test with no future_target only one step prediction test with all future data test for generation assert that the last element of the lagged sequence is the one after the encoders input assert that the first element of the future_values is offset by lag after the decoders input
import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download
from parameterized import parameterized

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 (
        TimeSeriesTransformerConfig,
        TimeSeriesTransformerForPrediction,
        TimeSeriesTransformerModel,
    )
    from transformers.models.time_series_transformer.modeling_time_series_transformer import (
        TimeSeriesTransformerDecoder,
        TimeSeriesTransformerEncoder,
    )


@require_torch
class TimeSeriesTransformerModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        prediction_length=7,
        context_length=14,
        cardinality=19,
        embedding_dimension=5,
        num_time_features=4,
        is_training=True,
        hidden_size=64,
        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],
    ):
        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

    def get_config(self):
        return TimeSeriesTransformerConfig(
            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,
            lags_sequence=self.lags_sequence,
            num_time_features=self.num_time_features,
            num_static_real_features=1,
            num_static_categorical_features=1,
            cardinality=[self.cardinality],
            embedding_dimension=[self.embedding_dimension],
        )

    def prepare_time_series_transformer_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])
        static_real_features = floats_tensor([self.batch_size, 1])

        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,
            "static_real_features": static_real_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_time_series_transformer_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 = TimeSeriesTransformerModel(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 = TimeSeriesTransformerEncoder.from_pretrained(tmpdirname).to(torch_device)

        transformer_inputs, _, _, _ = model.create_network_inputs(**inputs_dict)
        enc_input = transformer_inputs[:, : config.context_length, ...]
        dec_input = transformer_inputs[:, config.context_length :, ...]

        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)

        with tempfile.TemporaryDirectory() as tmpdirname:
            decoder = model.get_decoder()
            decoder.save_pretrained(tmpdirname)
            decoder = TimeSeriesTransformerDecoder.from_pretrained(tmpdirname).to(torch_device)

        last_hidden_state_2 = decoder(
            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 TimeSeriesTransformerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (
        (TimeSeriesTransformerModel, TimeSeriesTransformerForPrediction) if is_torch_available() else ()
    )
    all_generative_model_classes = (TimeSeriesTransformerForPrediction,) if is_torch_available() else ()
    pipeline_model_mapping = {"feature-extraction": TimeSeriesTransformerModel} if is_torch_available() else {}
    is_encoder_decoder = True
    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 = TimeSeriesTransformerModelTester(self)
        self.config_tester = ConfigTester(
            self,
            config_class=TimeSeriesTransformerConfig,
            has_text_modality=False,
            prediction_length=self.model_tester.prediction_length,
        )

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

    def test_save_load_strict(self):
        config, _ = 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)

    
    def test_resize_tokens_embeddings(self):
        pass

    
    def test_model_main_input_name(self):
        model_signature = inspect.signature(getattr(TimeSeriesTransformerModel, "forward"))
        
        observed_main_input_name = list(model_signature.parameters.keys())[1]
        self.assertEqual(TimeSeriesTransformerModel.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",
            ]

            expected_arg_names.extend(
                [
                    "future_observed_mask",
                    "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",
                ]
                if "future_observed_mask" in arg_names
                else [
                    "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)

        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, encoder_seq_length],
            )
            out_len = len(outputs)

            correct_outlen = 7

            if "last_hidden_state" 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, decoder_seq_length],
            )

            
            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,
                    encoder_seq_length,
                ],
            )

        
        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, encoder_seq_length],
        )

    @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(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

    @parameterized.expand(
        [
            (1, 5, [1]),
            (1, 5, [1, 10, 15]),
            (1, 5, [3, 6, 9, 10]),
            (2, 5, [1, 2, 7]),
            (2, 5, [2, 3, 4, 6]),
            (4, 5, [1, 5, 9, 11]),
            (4, 5, [7, 8, 13, 14]),
        ],
    )
    def test_create_network_inputs(self, prediction_length, context_length, lags_sequence):
        history_length = max(lags_sequence) + context_length

        config = TimeSeriesTransformerConfig(
            prediction_length=prediction_length,
            context_length=context_length,
            lags_sequence=lags_sequence,
            scaling=False,
            num_parallel_samples=10,
            num_static_categorical_features=1,
            cardinality=[1],
            embedding_dimension=[2],
            num_static_real_features=1,
        )
        model = TimeSeriesTransformerModel(config)

        batch = {
            "static_categorical_features": torch.tensor([[0]], dtype=torch.int64),
            "static_real_features": torch.tensor([[0.0]], dtype=torch.float32),
            "past_time_features": torch.arange(history_length, dtype=torch.float32).view(1, history_length, 1),
            "past_values": torch.arange(history_length, dtype=torch.float32).view(1, history_length),
            "past_observed_mask": torch.arange(history_length, dtype=torch.float32).view(1, history_length),
        }

        
        batch["future_time_features"] = torch.arange(history_length, history_length + 1, dtype=torch.float32).view(
            1, 1, 1
        )
        transformer_inputs, loc, scale, _ = model.create_network_inputs(**batch)

        self.assertTrue((scale == 1.0).all())
        assert (loc == 0.0).all()

        ref = torch.arange(max(lags_sequence), history_length, dtype=torch.float32)

        for idx, lag in enumerate(lags_sequence):
            assert torch.isclose(ref - lag, transformer_inputs[0, :, idx]).all()

        
        batch["future_time_features"] = torch.arange(
            history_length, history_length + prediction_length, dtype=torch.float32
        ).view(1, prediction_length, 1)
        batch["future_values"] = torch.arange(
            history_length, history_length + prediction_length, dtype=torch.float32
        ).view(1, prediction_length)
        transformer_inputs, loc, scale, _ = model.create_network_inputs(**batch)

        assert (scale == 1.0).all()
        assert (loc == 0.0).all()

        ref = torch.arange(max(lags_sequence), history_length + prediction_length, dtype=torch.float32)

        for idx, lag in enumerate(lags_sequence):
            assert torch.isclose(ref - lag, transformer_inputs[0, :, idx]).all()

        
        batch.pop("future_values")
        transformer_inputs, loc, scale, _ = model.create_network_inputs(**batch)

        lagged_sequence = model.get_lagged_subsequences(
            sequence=batch["past_values"],
            subsequences_length=1,
            shift=1,
        )
        
        assert transformer_inputs[0, ..., 0][-1] + 1 == lagged_sequence[0, ..., 0][-1]

        future_values = torch.arange(history_length, history_length + prediction_length, dtype=torch.float32).view(
            1, prediction_length
        )
        
        assert lagged_sequence[0, ..., 0][-1] + lags_sequence[0] == future_values[0, ..., 0]

    @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 TimeSeriesTransformerModelIntegrationTests(unittest.TestCase):
    def test_inference_no_head(self):
        model = TimeSeriesTransformerModel.from_pretrained("huggingface/time-series-transformer-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"],
                static_real_features=batch["static_real_features"],
                future_values=batch["future_values"],
                future_time_features=batch["future_time_features"],
            ).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.8196, -1.5131, 1.4620], [1.1268, -1.3238, 1.5997], [1.5098, -1.0715, 1.7359]], device=torch_device
        )
        self.assertTrue(torch.allclose(output[0, :3, :3], expected_slice, atol=TOLERANCE))

    def test_inference_head(self):
        model = TimeSeriesTransformerForPrediction.from_pretrained(
            "huggingface/time-series-transformer-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"],
                static_real_features=batch["static_real_features"],
                future_time_features=batch["future_time_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(
            [[-1.2957, -1.0280, -0.6045], [-0.7017, -0.8193, -0.3717], [-1.0449, -0.8149, 0.1405]], device=torch_device
        )
        self.assertTrue(torch.allclose(output[0, :3, :3], expected_slice, atol=TOLERANCE))

    def test_seq_to_seq_generation(self):
        model = TimeSeriesTransformerForPrediction.from_pretrained(
            "huggingface/time-series-transformer-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"],
                static_real_features=batch["static_real_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([2825.2749, 3584.9207, 6763.9951], device=torch_device)
        mean_prediction = outputs.sequences.mean(dim=1)
        self.assertTrue(torch.allclose(mean_prediction[0, -3:], expected_slice, rtol=1e-1))