codingutf8 2022 the huggingface inc team 2018 amazon com inc or its affiliates 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 pytorch time series transformer model from typing import list optional tuple union import numpy as np import torch from torch import nn from activations import act2fn from modelingattnmaskutils import prepare4dattentionmask prepare4dcausalattentionmask from modelingoutputs import basemodeloutput basemodeloutputwithpastandcrossattentions sampletspredictionoutput seq2seqtsmodeloutput seq2seqtspredictionoutput from modelingutils import pretrainedmodel from timeseriesutils import negativebinomialoutput normaloutput studenttoutput from utils import addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationtimeseriestransformer import timeseriestransformerconfig logger logging getloggername configfordoc timeseriestransformerconfig timeseriestransformerpretrainedmodelarchivelist huggingfacetimeseriestransformertourismmonthly see all timeseriestransformer models at https huggingface comodels filtertimeseriestransformer class timeseriesfeatureembeddernn module def initself cardinalities listint embeddingdims listint none super init self numfeatures lencardinalities self embedders nn modulelistnn embeddingc d for c d in zipcardinalities embeddingdims def forwardself features torch tensor torch tensor if self numfeatures 1 we slice the last dimension giving an array of length self numfeatures with shape n t or n catfeatureslices torch chunkfeatures self numfeatures dim1 else catfeatureslices features return torch cat embedcatfeatureslice squeeze1 for embed catfeatureslice in zipself embedders catfeatureslices dim1 class timeseriesstdscalernn module def initself config timeseriestransformerconfig super init self dim config scalingdim if hasattrconfig scalingdim else 1 self keepdim config keepdim if hasattrconfig keepdim else true self minimumscale config minimumscale if hasattrconfig minimumscale else 1e5 def forward self data torch tensor observedindicator torch tensor tupletorch tensor torch tensor torch tensor denominator observedindicator sumself dim keepdimself keepdim denominator denominator clampmin1 0 loc data observedindicator sumself dim keepdimself keepdim denominator variance data loc observedindicator 2 sumself dim keepdimself keepdim denominator scale torch sqrtvariance self minimumscale return data loc scale loc scale class timeseriesmeanscalernn module def initself config timeseriestransformerconfig super init self dim config scalingdim if hasattrconfig scalingdim else 1 self keepdim config keepdim if hasattrconfig keepdim else true self minimumscale config minimumscale if hasattrconfig minimumscale else 1e10 self defaultscale config defaultscale if hasattrconfig defaultscale else none def forward self data torch tensor observedindicator torch tensor tupletorch tensor torch tensor torch tensor tssum data observedindicator abs sumself dim keepdimtrue numobserved observedindicator sumself dim keepdimtrue scale tssum torch clampnumobserved min1 if defaultscale is provided we use it otherwise we use the scale of the batch if self defaultscale is none batchsum tssum sumdim0 batchobservations torch clampnumobserved sum0 min1 defaultscale torch squeezebatchsum batchobservations else defaultscale self defaultscale torch oneslikescale apply default scale where there are no observations scale torch wherenumobserved 0 scale defaultscale ensure the scale is at least self minimumscale scale torch clampscale minself minimumscale scaleddata data scale if not self keepdim scale scale squeezedimself dim return scaleddata torch zeroslikescale scale class timeseriesnopscalernn module def initself config timeseriestransformerconfig super init self dim config scalingdim if hasattrconfig scalingdim else 1 self keepdim config keepdim if hasattrconfig keepdim else true def forward self data torch tensor observedindicator torch tensor none tupletorch tensor torch tensor torch tensor scale torch oneslikedata requiresgradfalse meandimself dim keepdimself keepdim loc torch zeroslikedata requiresgradfalse meandimself dim keepdimself keepdim return data loc scale def nllinput torch distributions distribution target torch tensor torch tensor return input logprobtarget def weightedaverageinputtensor torch tensor weights optionaltorch tensor none dimnone torch tensor if weights is not none weightedtensor torch whereweights 0 inputtensor weights torch zeroslikeinputtensor sumweights torch clampweights sumdimdim if dim else weights sum min1 0 return weightedtensor sumdimdim if dim else weightedtensor sum sumweights else return inputtensor meandimdim copied from transformers models marian modelingmarian mariansinusoidalpositionalembedding with mariantimeseries class timeseriessinusoidalpositionalembeddingnn embedding identical to the xlm createsinusoidalembeddings except features are not interleaved the cos features are in the 2nd half of the vector dim 2 inputidsshape is expected to be bsz x seqlen bsz seqlen inputidsshape 2 positions torch arange pastkeyvalueslength pastkeyvalueslength seqlen dtypetorch long deviceself weight device return super forwardpositions class timeseriesvalueembeddingnn module def initself featuresize dmodel super init self valueprojection nn linearinfeaturesfeaturesize outfeaturesdmodel biasfalse def forwardself x return self valueprojectionx copied from transformers models bart modelingbart bartattention with barttimeseriestransformer class timeseriestransformerattentionnn module input shape batch x time x channel if keyvaluestates are provided this layer is used as a crossattention layer for the decoder iscrossattention keyvaluestates is not none bsz tgtlen hiddenstates size get query proj querystates self qprojhiddenstates self scaling get key value proj pastkeyvalue0 shape2 keyvaluestates shape1 is checking that the sequencelength of the pastkeyvalue is the same as the provided keyvaluestates to support prefix tuning if iscrossattention and pastkeyvalue is not none and pastkeyvalue0 shape2 keyvaluestates shape1 reuse k v crossattentions keystates pastkeyvalue0 valuestates pastkeyvalue1 elif iscrossattention crossattentions keystates self shapeself kprojkeyvaluestates 1 bsz valuestates self shapeself vprojkeyvaluestates 1 bsz elif pastkeyvalue is not none reuse k v selfattention keystates self shapeself kprojhiddenstates 1 bsz valuestates self shapeself vprojhiddenstates 1 bsz keystates torch catpastkeyvalue0 keystates dim2 valuestates torch catpastkeyvalue1 valuestates dim2 else selfattention keystates self shapeself kprojhiddenstates 1 bsz valuestates self shapeself vprojhiddenstates 1 bsz if self isdecoder if crossattention save tupletorch tensor torch tensor of all cross attention keyvaluestates further calls to crossattention layer can then reuse all crossattention keyvaluestates first if case if unidirectional selfattention decoder save tupletorch tensor torch tensor of all previous decoder keyvaluestates further calls to unidirectional selfattention can concat previous decoder keyvaluestates to current projected keyvaluestates third elif case if encoder bidirectional selfattention pastkeyvalue is always none pastkeyvalue keystates valuestates projshape bsz self numheads 1 self headdim querystates self shapequerystates tgtlen bsz viewprojshape keystates keystates reshapeprojshape valuestates valuestates reshapeprojshape srclen keystates size1 attnweights torch bmmquerystates keystates transpose1 2 if attnweights size bsz self numheads tgtlen srclen raise valueerror fattention weights should be of size bsz self numheads tgtlen srclen but is f attnweights size if attentionmask is not none if attentionmask size bsz 1 tgtlen srclen raise valueerror fattention mask should be of size bsz 1 tgtlen srclen but is attentionmask size attnweights attnweights viewbsz self numheads tgtlen srclen attentionmask attnweights attnweights viewbsz self numheads tgtlen srclen attnweights nn functional softmaxattnweights dim1 if layerheadmask is not none if layerheadmask size self numheads raise valueerror fhead mask for a single layer should be of size self numheads but is f layerheadmask size attnweights layerheadmask view1 1 1 1 attnweights viewbsz self numheads tgtlen srclen attnweights attnweights viewbsz self numheads tgtlen srclen if outputattentions this operation is a bit awkward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to be reshaped twice and have to be reused in the following attnweightsreshaped attnweights viewbsz self numheads tgtlen srclen attnweights attnweightsreshaped viewbsz self numheads tgtlen srclen else attnweightsreshaped none attnprobs nn functional dropoutattnweights pself dropout trainingself training attnoutput torch bmmattnprobs valuestates if attnoutput size bsz self numheads tgtlen self headdim raise valueerror fattnoutput should be of size bsz self numheads tgtlen self headdim but is f attnoutput size attnoutput attnoutput viewbsz self numheads tgtlen self headdim attnoutput attnoutput transpose1 2 use the embeddim from the config stored in the class rather than hiddenstate because attnoutput can be partitioned across gpus when using tensorparallelism attnoutput attnoutput reshapebsz tgtlen self embeddim attnoutput self outprojattnoutput return attnoutput attnweightsreshaped pastkeyvalue copied from transformers models bart modelingbart bartencoderlayer with barttimeseriestransformer barttimeseriestransformer class timeseriestransformerencoderlayernn module def initself config timeseriestransformerconfig super init self embeddim config dmodel attntype flashattention2 if getattrconfig flashattn2enabled false else default self selfattn timeseriestransformerattentionclassesattntype embeddimself embeddim numheadsconfig encoderattentionheads dropoutconfig attentiondropout configconfig self selfattnlayernorm nn layernormself embeddim self dropout config dropout self activationfn act2fnconfig activationfunction self activationdropout config activationdropout self fc1 nn linearself embeddim config encoderffndim self fc2 nn linearconfig encoderffndim self embeddim self finallayernorm nn layernormself embeddim def forward self hiddenstates torch floattensor attentionmask torch floattensor layerheadmask torch floattensor outputattentions optionalbool false tupletorch floattensor optionaltorch floattensor residual hiddenstates hiddenstates attnweights self selfattn hiddenstateshiddenstates attentionmaskattentionmask layerheadmasklayerheadmask outputattentionsoutputattentions hiddenstates nn functional dropouthiddenstates pself dropout trainingself training hiddenstates residual hiddenstates hiddenstates self selfattnlayernormhiddenstates residual hiddenstates hiddenstates self activationfnself fc1hiddenstates hiddenstates nn functional dropouthiddenstates pself activationdropout trainingself training hiddenstates self fc2hiddenstates hiddenstates nn functional dropouthiddenstates pself dropout trainingself training hiddenstates residual hiddenstates hiddenstates self finallayernormhiddenstates if hiddenstates dtype torch float16 and torch isinfhiddenstates any or torch isnanhiddenstates any clampvalue torch finfohiddenstates dtype max 1000 hiddenstates torch clamphiddenstates minclampvalue maxclampvalue outputs hiddenstates if outputattentions outputs attnweights return outputs timeseriestransformerattentionclasses default timeseriestransformerattention copied from transformers models bart modelingbart bartdecoderlayer with barttimeseriestransformer with barttimeseriestransformer class timeseriestransformerdecoderlayernn module def initself config timeseriestransformerconfig super init self embeddim config dmodel attntype flashattention2 if getattrconfig flashattn2enabled false else default self selfattn timeseriestransformerattentionclassesattntype embeddimself embeddim numheadsconfig decoderattentionheads dropoutconfig attentiondropout isdecodertrue iscausaltrue configconfig self dropout config dropout self activationfn act2fnconfig activationfunction self activationdropout config activationdropout self selfattnlayernorm nn layernormself embeddim self encoderattn timeseriestransformerattentionclassesattntype self embeddim config decoderattentionheads dropoutconfig attentiondropout isdecodertrue configconfig self encoderattnlayernorm nn layernormself embeddim self fc1 nn linearself embeddim config decoderffndim self fc2 nn linearconfig decoderffndim self embeddim self finallayernorm nn layernormself embeddim def forward self hiddenstates torch tensor attentionmask optionaltorch tensor none encoderhiddenstates optionaltorch tensor none encoderattentionmask optionaltorch tensor none layerheadmask optionaltorch tensor none crossattnlayerheadmask optionaltorch tensor none pastkeyvalue optionaltupletorch tensor none outputattentions optionalbool false usecache optionalbool true tupletorch floattensor optionaltupletorch floattensor torch floattensor residual hiddenstates self attention decoder unidirectional selfattention cached keyvalues tuple is at positions 1 2 selfattnpastkeyvalue pastkeyvalue 2 if pastkeyvalue is not none else none add present selfattn cache to positions 1 2 of presentkeyvalue tuple hiddenstates selfattnweights presentkeyvalue self selfattn hiddenstateshiddenstates pastkeyvalueselfattnpastkeyvalue attentionmaskattentionmask layerheadmasklayerheadmask outputattentionsoutputattentions hiddenstates nn functional dropouthiddenstates pself dropout trainingself training hiddenstates residual hiddenstates hiddenstates self selfattnlayernormhiddenstates crossattention block crossattnpresentkeyvalue none crossattnweights none if encoderhiddenstates is not none residual hiddenstates crossattn cached keyvalues tuple is at positions 3 4 of presentkeyvalue tuple crossattnpastkeyvalue pastkeyvalue2 if pastkeyvalue is not none else none hiddenstates crossattnweights crossattnpresentkeyvalue self encoderattn hiddenstateshiddenstates keyvaluestatesencoderhiddenstates attentionmaskencoderattentionmask layerheadmaskcrossattnlayerheadmask pastkeyvaluecrossattnpastkeyvalue outputattentionsoutputattentions hiddenstates nn functional dropouthiddenstates pself dropout trainingself training hiddenstates residual hiddenstates hiddenstates self encoderattnlayernormhiddenstates add crossattn to positions 3 4 of presentkeyvalue tuple presentkeyvalue presentkeyvalue crossattnpresentkeyvalue fully connected residual hiddenstates hiddenstates self activationfnself fc1hiddenstates hiddenstates nn functional dropouthiddenstates pself activationdropout trainingself training hiddenstates self fc2hiddenstates hiddenstates nn functional dropouthiddenstates pself dropout trainingself training hiddenstates residual hiddenstates hiddenstates self finallayernormhiddenstates outputs hiddenstates if outputattentions outputs selfattnweights crossattnweights if usecache outputs presentkeyvalue return outputs class timeseriestransformerpretrainedmodelpretrainedmodel configclass timeseriestransformerconfig basemodelprefix model maininputname pastvalues supportsgradientcheckpointing true def initweightsself module std self config initstd if isinstancemodule nn linear module weight data normalmean0 0 stdstd if module bias is not none module bias data zero elif isinstancemodule timeseriessinusoidalpositionalembedding pass elif isinstancemodule nn embedding module weight data normalmean0 0 stdstd if module paddingidx is not none module weight datamodule paddingidx zero timeseriestransformerstartdocstring r this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a pytorch torch nn modulehttps pytorch orgdocsstablenn htmltorch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config timeseriestransformerconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel frompretrained method to load the model weights class timeseriestransformerencodertimeseriestransformerpretrainedmodel def initself config timeseriestransformerconfig super initconfig self dropout config dropout self layerdrop config encoderlayerdrop if config predictionlength is none raise valueerrorthe predictionlength config needs to be specified self valueembedding timeseriesvalueembeddingfeaturesizeconfig featuresize dmodelconfig dmodel self embedpositions timeseriessinusoidalpositionalembedding config contextlength config predictionlength config dmodel self layers nn modulelisttimeseriestransformerencoderlayerconfig for in rangeconfig encoderlayers self layernormembedding nn layernormconfig dmodel self gradientcheckpointing false initialize weights and apply final processing self postinit def forward self attentionmask optionaltorch tensor none headmask optionaltorch tensor none inputsembeds optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutput r args attentionmask torch tensor of shape batchsize sequencelength optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask headmask torch tensor of shape encoderlayers encoderattentionheads optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked inputsembeds torch floattensor of shape batchsize sequencelength hiddensize optional optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple expand attentionmask bsz seqlen bsz 1 tgtseqlen srcseqlen check if headmask has a correct number of layers specified if desired add layerdrop see https arxiv orgabs1909 11556 for description transformer decoder consisting of config decoderlayers layers each layer is a timeseriestransformerdecoderlayer args config timeseriestransformerconfig initialize weights and apply final processing outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates usecache usecache if usecache is not none else self config usecache returndict returndict if returndict is not none else self config usereturndict inputshape inputsembeds size 1 pastkeyvalueslength pastkeyvalueslength pastkeyvalues00 shape2 if pastkeyvalues is not none else 0 attentionmask prepare4dcausalattentionmask attentionmask inputshape inputsembeds pastkeyvalueslength expand encoder attention mask if encoderhiddenstates is not none and encoderattentionmask is not none bsz seqlen bsz 1 tgtseqlen srcseqlen encoderattentionmask prepare4dattentionmask encoderattentionmask inputsembeds dtype tgtleninputshape1 hiddenstates self valueembeddinginputsembeds embedpos self embedpositionsinputsembeds size pastkeyvalueslengthself config contextlength hiddenstates self layernormembeddinghiddenstates embedpos hiddenstates nn functional dropouthiddenstates pself dropout trainingself training if self gradientcheckpointing and self training if usecache logger warningonce usecachetrue is incompatible with gradient checkpointing setting usecachefalse usecache false decoder layers allhiddenstates if outputhiddenstates else none allselfattns if outputattentions else none allcrossattentions if outputattentions and encoderhiddenstates is not none else none nextdecodercache if usecache else none check if headmaskcrossattnheadmask has a correct number of layers specified if desired for attnmask maskname in zipheadmask crossattnheadmask headmask crossattnheadmask if attnmask is not none if attnmask size0 lenself layers raise valueerror fthe maskname should be specified for lenself layers layers but it is for f headmask size0 for idx decoderlayer in enumerateself layers add layerdrop see https arxiv orgabs1909 11556 for description if outputhiddenstates allhiddenstates hiddenstates if self training dropoutprobability torch rand if dropoutprobability self layerdrop continue pastkeyvalue pastkeyvaluesidx if pastkeyvalues is not none else none if self gradientcheckpointing and self training layeroutputs self gradientcheckpointingfunc decoderlayer call hiddenstates attentionmask encoderhiddenstates encoderattentionmask headmaskidx if headmask is not none else none crossattnheadmaskidx if crossattnheadmask is not none else none none outputattentions usecache else layeroutputs decoderlayer hiddenstates attentionmaskattentionmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask layerheadmaskheadmaskidx if headmask is not none else none crossattnlayerheadmask crossattnheadmaskidx if crossattnheadmask is not none else none pastkeyvaluepastkeyvalue outputattentionsoutputattentions usecacheusecache hiddenstates layeroutputs0 if usecache nextdecodercache layeroutputs3 if outputattentions else 1 if outputattentions allselfattns layeroutputs1 if encoderhiddenstates is not none allcrossattentions layeroutputs2 add hidden states from the last decoder layer if outputhiddenstates allhiddenstates hiddenstates nextcache nextdecodercache if usecache else none if not returndict return tuple v for v in hiddenstates nextcache allhiddenstates allselfattns allcrossattentions if v is not none return basemodeloutputwithpastandcrossattentions lasthiddenstatehiddenstates pastkeyvaluesnextcache hiddenstatesallhiddenstates attentionsallselfattns crossattentionsallcrossattentions addstartdocstrings the bare time series transformer model outputting raw hiddenstates without any specific head on top timeseriestransformerstartdocstring class timeseriestransformermodeltimeseriestransformerpretrainedmodel def initself config timeseriestransformerconfig super initconfig if config scaling mean or config scaling is true self scaler timeseriesmeanscalerconfig elif config scaling std self scaler timeseriesstdscalerconfig else self scaler timeseriesnopscalerconfig if config numstaticcategoricalfeatures 0 self embedder timeseriesfeatureembedder cardinalitiesconfig cardinality embeddingdimsconfig embeddingdimension transformer encoderdecoder and mask initializer self encoder timeseriestransformerencoderconfig self decoder timeseriestransformerdecoderconfig initialize weights and apply final processing self postinit property def pastlengthself int return self config contextlength maxself config lagssequence def getlaggedsubsequences self sequence torch tensor subsequenceslength int shift int 0 torch tensor sequencelength sequence shape1 indices lag shift for lag in self config lagssequence if maxindices subsequenceslength sequencelength raise valueerror flags cannot go further than history length found lag maxindices fwhile history length is only sequencelength laggedvalues for lagindex in indices beginindex lagindex subsequenceslength endindex lagindex if lagindex 0 else none laggedvalues appendsequence beginindex endindex return torch stacklaggedvalues dim1 def createnetworkinputs self pastvalues torch tensor pasttimefeatures torch tensor staticcategoricalfeatures optionaltorch tensor none staticrealfeatures optionaltorch tensor none pastobservedmask optionaltorch tensor none futurevalues optionaltorch tensor none futuretimefeatures optionaltorch tensor none time feature timefeat torch cat pasttimefeatures self pastlength self config contextlength futuretimefeatures dim1 if futurevalues is not none else pasttimefeatures self pastlength self config contextlength target if pastobservedmask is none pastobservedmask torch oneslikepastvalues context pastvalues self config contextlength observedcontext pastobservedmask self config contextlength loc scale self scalercontext observedcontext inputs torch catpastvalues futurevalues dim1 loc scale if futurevalues is not none else pastvalues loc scale static features logabsloc loc abs log1p if self config inputsize 1 else loc squeeze1 abs log1p logscale scale log if self config inputsize 1 else scale squeeze1 log staticfeat torch catlogabsloc logscale dim1 if staticrealfeatures is not none staticfeat torch catstaticrealfeatures staticfeat dim1 if staticcategoricalfeatures is not none embeddedcat self embedderstaticcategoricalfeatures staticfeat torch catembeddedcat staticfeat dim1 expandedstaticfeat staticfeat unsqueeze1 expand1 timefeat shape1 1 all features features torch catexpandedstaticfeat timefeat dim1 lagged features subsequenceslength self config contextlength self config predictionlength if futurevalues is not none else self config contextlength laggedsequence self getlaggedsubsequencessequenceinputs subsequenceslengthsubsequenceslength lagsshape laggedsequence shape reshapedlaggedsequence laggedsequence reshapelagsshape0 lagsshape1 1 if reshapedlaggedsequence shape1 timefeat shape1 raise valueerror finput length reshapedlaggedsequence shape1 and time feature lengths timefeat shape1 does not match transformer inputs transformerinputs torch catreshapedlaggedsequence features dim1 return transformerinputs loc scale staticfeat def getencoderself return self encoder def getdecoderself return self decoder addstartdocstringstomodelforwardtimeseriestransformerinputsdocstring replacereturndocstringsoutputtypeseq2seqtsmodeloutput configclassconfigfordoc def forward self pastvalues torch tensor pasttimefeatures torch tensor pastobservedmask torch tensor staticcategoricalfeatures optionaltorch tensor none staticrealfeatures optionaltorch tensor none futurevalues optionaltorch tensor none futuretimefeatures optionaltorch tensor none decoderattentionmask optionaltorch longtensor none headmask optionaltorch tensor none decoderheadmask optionaltorch tensor none crossattnheadmask optionaltorch tensor none encoderoutputs optionallisttorch floattensor none pastkeyvalues optionallisttorch floattensor none outputhiddenstates optionalbool none outputattentions optionalbool none usecache optionalbool none returndict optionalbool none unionseq2seqtsmodeloutput tuple r returns examples python from huggingfacehub import hfhubdownload import torch from transformers import timeseriestransformermodel file hfhubdownload repoidhfinternaltestingtourismmonthlybatch filenametrainbatch pt repotypedataset batch torch loadfile model timeseriestransformermodel frompretrainedhuggingfacetimeseriestransformertourismmonthly during training one provides both past and future values as well as possible additional features outputs model pastvaluesbatchpastvalues pasttimefeaturesbatchpasttimefeatures pastobservedmaskbatchpastobservedmask staticcategoricalfeaturesbatchstaticcategoricalfeatures staticrealfeaturesbatchstaticrealfeatures futurevaluesbatchfuturevalues futuretimefeaturesbatchfuturetimefeatures lasthiddenstate outputs lasthiddenstate outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates usecache usecache if usecache is not none else self config usecache returndict returndict if returndict is not none else self config usereturndict transformerinputs loc scale staticfeat self createnetworkinputs pastvaluespastvalues pasttimefeaturespasttimefeatures pastobservedmaskpastobservedmask staticcategoricalfeaturesstaticcategoricalfeatures staticrealfeaturesstaticrealfeatures futurevaluesfuturevalues futuretimefeaturesfuturetimefeatures if encoderoutputs is none encinput transformerinputs self config contextlength encoderoutputs self encoder inputsembedsencinput headmaskheadmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict if the user passed a tuple for encoderoutputs we wrap it in a basemodeloutput when returndicttrue elif returndict and not isinstanceencoderoutputs basemodeloutput encoderoutputs basemodeloutput lasthiddenstateencoderoutputs0 hiddenstatesencoderoutputs1 if lenencoderoutputs 1 else none attentionsencoderoutputs2 if lenencoderoutputs 2 else none decinput transformerinputs self config contextlength decoderoutputs self decoder inputsembedsdecinput attentionmaskdecoderattentionmask encoderhiddenstatesencoderoutputs0 headmaskdecoderheadmask crossattnheadmaskcrossattnheadmask pastkeyvaluespastkeyvalues usecacheusecache outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict if not returndict return decoderoutputs encoderoutputs loc scale staticfeat return seq2seqtsmodeloutput lasthiddenstatedecoderoutputs lasthiddenstate pastkeyvaluesdecoderoutputs pastkeyvalues decoderhiddenstatesdecoderoutputs hiddenstates decoderattentionsdecoderoutputs attentions crossattentionsdecoderoutputs crossattentions encoderlasthiddenstateencoderoutputs lasthiddenstate encoderhiddenstatesencoderoutputs hiddenstates encoderattentionsencoderoutputs attentions locloc scalescale staticfeaturesstaticfeat addstartdocstrings the time series transformer model with a distribution head on top for timeseries forecasting timeseriestransformerstartdocstring class timeseriestransformerforpredictiontimeseriestransformerpretrainedmodel def initself config timeseriestransformerconfig super initconfig self model timeseriestransformermodelconfig if config distributionoutput studentt self distributionoutput studenttoutputdimconfig inputsize elif config distributionoutput normal self distributionoutput normaloutputdimconfig inputsize elif config distributionoutput negativebinomial self distributionoutput negativebinomialoutputdimconfig inputsize else raise valueerrorfunknown distribution output config distributionoutput self parameterprojection self distributionoutput getparameterprojectionself model config dmodel self targetshape self distributionoutput eventshape if config loss nll self loss nll else raise valueerrorfunknown loss function config loss initialize weights of distributionoutput and apply final processing self postinit def outputparamsself decoutput return self parameterprojectiondecoutput def getencoderself return self model getencoder def getdecoderself return self model getdecoder torch jit ignore def outputdistributionself params locnone scalenone trailingnnone torch distributions distribution slicedparams params if trailingn is not none slicedparams p trailingn for p in params return self distributionoutput distributionslicedparams locloc scalescale addstartdocstringstomodelforwardtimeseriestransformerinputsdocstring replacereturndocstringsoutputtypeseq2seqtsmodeloutput configclassconfigfordoc def forward self pastvalues torch tensor pasttimefeatures torch tensor pastobservedmask torch tensor staticcategoricalfeatures optionaltorch tensor none staticrealfeatures optionaltorch tensor none futurevalues optionaltorch tensor none futuretimefeatures optionaltorch tensor none futureobservedmask optionaltorch tensor none decoderattentionmask optionaltorch longtensor none headmask optionaltorch tensor none decoderheadmask optionaltorch tensor none crossattnheadmask optionaltorch tensor none encoderoutputs optionallisttorch floattensor none pastkeyvalues optionallisttorch floattensor none outputhiddenstates optionalbool none outputattentions optionalbool none usecache optionalbool none returndict optionalbool none unionseq2seqtsmodeloutput tuple r returns examples python from huggingfacehub import hfhubdownload import torch from transformers import timeseriestransformerforprediction file hfhubdownload repoidhfinternaltestingtourismmonthlybatch filenametrainbatch pt repotypedataset batch torch loadfile model timeseriestransformerforprediction frompretrained huggingfacetimeseriestransformertourismmonthly during training one provides both past and future values as well as possible additional features outputs model pastvaluesbatchpastvalues pasttimefeaturesbatchpasttimefeatures pastobservedmaskbatchpastobservedmask staticcategoricalfeaturesbatchstaticcategoricalfeatures staticrealfeaturesbatchstaticrealfeatures futurevaluesbatchfuturevalues futuretimefeaturesbatchfuturetimefeatures loss outputs loss loss backward during inference one only provides past values as well as possible additional features the model autoregressively generates future values outputs model generate pastvaluesbatchpastvalues pasttimefeaturesbatchpasttimefeatures pastobservedmaskbatchpastobservedmask staticcategoricalfeaturesbatchstaticcategoricalfeatures staticrealfeaturesbatchstaticrealfeatures futuretimefeaturesbatchfuturetimefeatures meanprediction outputs sequences meandim1 returndict returndict if returndict is not none else self config usereturndict if futurevalues is not none usecache false outputs self model pastvaluespastvalues pasttimefeaturespasttimefeatures pastobservedmaskpastobservedmask staticcategoricalfeaturesstaticcategoricalfeatures staticrealfeaturesstaticrealfeatures futurevaluesfuturevalues futuretimefeaturesfuturetimefeatures decoderattentionmaskdecoderattentionmask headmaskheadmask decoderheadmaskdecoderheadmask crossattnheadmaskcrossattnheadmask encoderoutputsencoderoutputs pastkeyvaluespastkeyvalues outputhiddenstatesoutputhiddenstates outputattentionsoutputattentions usecacheusecache returndictreturndict predictionloss none params none if futurevalues is not none params self outputparamsoutputs0 outputs lasthiddenstate loc is 3rd last and scale is 2nd last output distribution self outputdistributionparams locoutputs3 scaleoutputs2 loss self lossdistribution futurevalues if futureobservedmask is none futureobservedmask torch oneslikefuturevalues if lenself targetshape 0 lossweights futureobservedmask else lossweights futureobservedmask mindim1 keepdimfalse predictionloss weightedaverageloss weightslossweights if not returndict outputs params outputs1 if params is not none else outputs1 return predictionloss outputs if predictionloss is not none else outputs return seq2seqtspredictionoutput losspredictionloss paramsparams pastkeyvaluesoutputs pastkeyvalues decoderhiddenstatesoutputs decoderhiddenstates decoderattentionsoutputs decoderattentions crossattentionsoutputs crossattentions encoderlasthiddenstateoutputs encoderlasthiddenstate encoderhiddenstatesoutputs encoderhiddenstates encoderattentionsoutputs encoderattentions locoutputs loc scaleoutputs scale staticfeaturesoutputs staticfeatures torch nograd def generate self pastvalues torch tensor pasttimefeatures torch tensor futuretimefeatures torch tensor pastobservedmask optionaltorch tensor none staticcategoricalfeatures optionaltorch tensor none staticrealfeatures optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none sampletspredictionoutput r greedily generate sequences of sample predictions from a model with a probability distribution head parameters pastvalues torch floattensor of shape batchsize sequencelength or batchsize sequencelength inputsize past values of the time series that serve as context in order to predict the future the sequence size of this tensor must be larger than the contextlength of the model since the model will use the larger size to construct lag features i e additional values from the past which are added in order to serve as extra context the sequencelength here is equal to config contextlength maxconfig lagssequence which if no lagssequence is configured is equal to config contextlength 7 as by default the largest lookback index in config lagssequence is 7 the property pastlength returns the actual length of the past the pastvalues is what the transformer encoder gets as input with optional additional features such as staticcategoricalfeatures staticrealfeatures pasttimefeatures and lags optionally missing values need to be replaced with zeros and indicated via the pastobservedmask for multivariate time series the inputsize 1 dimension is required and corresponds to the number of variates in the time series per time step pasttimefeatures torch floattensor of shape batchsize sequencelength numfeatures required time features which the model internally will add to pastvalues these could be things like month of year day of the month etc encoded as vectors for instance as fourier features these could also be socalled age features which basically help the model know at which point in life a timeseries is age features have small values for distant past time steps and increase monotonically the more we approach the current time step holiday features are also a good example of time features these features serve as the positional encodings of the inputs so contrary to a model like bert where the position encodings are learned from scratch internally as parameters of the model the time series transformer requires to provide additional time features the time series transformer only learns additional embeddings for staticcategoricalfeatures additional dynamic real covariates can be concatenated to this tensor with the caveat that these features must but known at prediction time the numfeatures here is equal to config numtimefeatures config numdynamicrealfeatures futuretimefeatures torch floattensor of shape batchsize predictionlength numfeatures required time features for the prediction window which the model internally will add to sampled predictions these could be things like month of year day of the month etc encoded as vectors for instance as fourier features these could also be socalled age features which basically help the model know at which point in life a timeseries is age features have small values for distant past time steps and increase monotonically the more we approach the current time step holiday features are also a good example of time features these features serve as the positional encodings of the inputs so contrary to a model like bert where the position encodings are learned from scratch internally as parameters of the model the time series transformer requires to provide additional time features the time series transformer only learns additional embeddings for staticcategoricalfeatures additional dynamic real covariates can be concatenated to this tensor with the caveat that these features must but known at prediction time the numfeatures here is equal to config numtimefeatures config numdynamicrealfeatures pastobservedmask torch booltensor of shape batchsize sequencelength or batchsize sequencelength inputsize optional boolean mask to indicate which pastvalues were observed and which were missing mask values selected in 0 1 1 for values that are observed 0 for values that are missing i e nans that were replaced by zeros staticcategoricalfeatures torch longtensor of shape batchsize number of static categorical features optional optional static categorical features for which the model will learn an embedding which it will add to the values of the time series static categorical features are features which have the same value for all time steps static over time a typical example of a static categorical feature is a time series id staticrealfeatures torch floattensor of shape batchsize number of static real features optional optional static real features which the model will add to the values of the time series static real features are features which have the same value for all time steps static over time a typical example of a static real feature is promotion information outputattentions bool optional whether or not to return the attentions tensors of all attention layers outputhiddenstates bool optional whether or not to return the hidden states of all layers return sampletspredictionoutput where the outputs sequences tensor will have shape batchsize number of samples predictionlength or batchsize number of samples predictionlength inputsize for multivariate predictions greedy decoding coding utf 8 2022 the huggingface inc team 2018 amazon com inc or its affiliates 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 pytorch time series transformer model see all timeseriestransformer models at https huggingface co models filter time_series_transformer embed a sequence of categorical features args cardinalities list int list of cardinalities of the categorical features embedding_dims list int list of embedding dimensions of the categorical features we slice the last dimension giving an array of length self num_features with shape n t or n standardize features by calculating the mean and scaling along the first dimension and then normalizes it by subtracting from the mean and dividing by the standard deviation parameters data torch tensor of shape batch_size sequence_length num_input_channels input for batch norm calculation observed_indicator torch booltensor of shape batch_size sequence_length num_input_channels calculating the scale on the observed indicator returns tuple of torch tensor of shapes batch_size sequence_length num_input_channels batch_size 1 num_input_channels batch_size 1 num_input_channels computes a scaling factor as the weighted average absolute value along the first dimension and scales the data accordingly parameters data torch tensor of shape batch_size sequence_length num_input_channels input for batch norm calculation observed_indicator torch booltensor of shape batch_size sequence_length num_input_channels calculating the scale on the observed indicator returns tuple of torch tensor of shapes batch_size sequence_length num_input_channels batch_size 1 num_input_channels batch_size 1 num_input_channels if default_scale is provided we use it otherwise we use the scale of the batch apply default scale where there are no observations ensure the scale is at least self minimum_scale assigns a scaling factor equal to 1 along the first dimension and therefore applies no scaling to the input data parameters data torch tensor of shape batch_size sequence_length num_input_channels input for batch norm calculation returns tuple of torch tensor of shapes batch_size sequence_length num_input_channels batch_size 1 num_input_channels batch_size 1 num_input_channels computes the negative log likelihood loss from input distribution with respect to target computes the weighted average of a given tensor across a given dim masking values associated with weight zero meaning instead of nan 0 nan you will get 0 0 0 args input_tensor torch floattensor input tensor of which the average must be computed weights torch floattensor optional weights tensor of the same shape as input_tensor dim int optional the dim along which to average input_tensor returns torch floattensor the tensor with values averaged along the specified dim copied from transformers models marian modeling_marian mariansinusoidalpositionalembedding with marian timeseries this module produces sinusoidal positional embeddings of any length identical to the xlm create_sinusoidal_embeddings except features are not interleaved the cos features are in the 2nd half of the vector dim 2 set early to avoid an error in pytorch 1 8 input_ids_shape is expected to be bsz x seqlen copied from transformers models bart modeling_bart bartattention with bart timeseriestransformer multi headed attention from attention is all you need paper input shape batch x time x channel if key_value_states are provided this layer is used as a cross attention layer for the decoder get query proj get key value proj past_key_value 0 shape 2 key_value_states shape 1 is checking that the sequence_length of the past_key_value is the same as the provided key_value_states to support prefix tuning reuse k v cross_attentions cross_attentions reuse k v self_attention self_attention if cross_attention save tuple torch tensor torch tensor of all cross attention key value_states further calls to cross_attention layer can then reuse all cross attention key value_states first if case if uni directional self attention decoder save tuple torch tensor torch tensor of all previous decoder key value_states further calls to uni directional self attention can concat previous decoder key value_states to current projected key value_states third elif case if encoder bi directional self attention past_key_value is always none this operation is a bit awkward but it s required to make sure that attn_weights keeps its gradient in order to do so attn_weights have to be reshaped twice and have to be reused in the following use the embed_dim from the config stored in the class rather than hidden_state because attn_output can be partitioned across gpus when using tensor parallelism copied from transformers models bart modeling_bart bartencoderlayer with bart timeseriestransformer bart time_series_transformer args hidden_states torch floattensor input to the layer of shape batch seq_len embed_dim attention_mask torch floattensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values layer_head_mask torch floattensor mask for attention heads in a given layer of size encoder_attention_heads output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail copied from transformers models bart modeling_bart bartdecoderlayer with bart timeseriestransformer with bart time_series_transformer args hidden_states torch floattensor input to the layer of shape batch seq_len embed_dim attention_mask torch floattensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values encoder_hidden_states torch floattensor cross attention input to the layer of shape batch seq_len embed_dim encoder_attention_mask torch floattensor encoder attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values layer_head_mask torch floattensor mask for attention heads in a given layer of size encoder_attention_heads cross_attn_layer_head_mask torch floattensor mask for cross attention heads in a given layer of size decoder_attention_heads past_key_value tuple torch floattensor cached past key and value projection states output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail self attention decoder uni directional self attention cached key values tuple is at positions 1 2 add present self attn cache to positions 1 2 of present_key_value tuple cross attention block cross_attn cached key values tuple is at positions 3 4 of present_key_value tuple add cross attn to positions 3 4 of present_key_value tuple fully connected this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config timeseriestransformerconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the pretrainedmodel from_pretrained method to load the model weights args past_values torch floattensor of shape batch_size sequence_length or batch_size sequence_length input_size past values of the time series that serve as context in order to predict the future the sequence size of this tensor must be larger than the context_length of the model since the model will use the larger size to construct lag features i e additional values from the past which are added in order to serve as extra context the sequence_length here is equal to config context_length max config lags_sequence which if no lags_sequence is configured is equal to config context_length 7 as by default the largest look back index in config lags_sequence is 7 the property _past_length returns the actual length of the past the past_values is what the transformer encoder gets as input with optional additional features such as static_categorical_features static_real_features past_time_features and lags optionally missing values need to be replaced with zeros and indicated via the past_observed_mask for multivariate time series the input_size 1 dimension is required and corresponds to the number of variates in the time series per time step past_time_features torch floattensor of shape batch_size sequence_length num_features required time features which the model internally will add to past_values these could be things like month of year day of the month etc encoded as vectors for instance as fourier features these could also be so called age features which basically help the model know at which point in life a time series is age features have small values for distant past time steps and increase monotonically the more we approach the current time step holiday features are also a good example of time features these features serve as the positional encodings of the inputs so contrary to a model like bert where the position encodings are learned from scratch internally as parameters of the model the time series transformer requires to provide additional time features the time series transformer only learns additional embeddings for static_categorical_features additional dynamic real covariates can be concatenated to this tensor with the caveat that these features must but known at prediction time the num_features here is equal to config num_time_features config num_dynamic_real_features past_observed_mask torch booltensor of shape batch_size sequence_length or batch_size sequence_length input_size optional boolean mask to indicate which past_values were observed and which were missing mask values selected in 0 1 1 for values that are observed 0 for values that are missing i e nans that were replaced by zeros static_categorical_features torch longtensor of shape batch_size number of static categorical features optional optional static categorical features for which the model will learn an embedding which it will add to the values of the time series static categorical features are features which have the same value for all time steps static over time a typical example of a static categorical feature is a time series id static_real_features torch floattensor of shape batch_size number of static real features optional optional static real features which the model will add to the values of the time series static real features are features which have the same value for all time steps static over time a typical example of a static real feature is promotion information future_values torch floattensor of shape batch_size prediction_length or batch_size prediction_length input_size optional future values of the time series that serve as labels for the model the future_values is what the transformer needs during training to learn to output given the past_values the sequence length here is equal to prediction_length see the demo notebook and code snippets for details optionally during training any missing values need to be replaced with zeros and indicated via the future_observed_mask for multivariate time series the input_size 1 dimension is required and corresponds to the number of variates in the time series per time step future_time_features torch floattensor of shape batch_size prediction_length num_features required time features for the prediction window which the model internally will add to future_values these could be things like month of year day of the month etc encoded as vectors for instance as fourier features these could also be so called age features which basically help the model know at which point in life a time series is age features have small values for distant past time steps and increase monotonically the more we approach the current time step holiday features are also a good example of time features these features serve as the positional encodings of the inputs so contrary to a model like bert where the position encodings are learned from scratch internally as parameters of the model the time series transformer requires to provide additional time features the time series transformer only learns additional embeddings for static_categorical_features additional dynamic real covariates can be concatenated to this tensor with the caveat that these features must but known at prediction time the num_features here is equal to config num_time_features config num_dynamic_real_features future_observed_mask torch booltensor of shape batch_size sequence_length or batch_size sequence_length input_size optional boolean mask to indicate which future_values were observed and which were missing mask values selected in 0 1 1 for values that are observed 0 for values that are missing i e nans that were replaced by zeros this mask is used to filter out missing values for the final loss calculation attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on certain token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask decoder_attention_mask torch longtensor of shape batch_size target_sequence_length optional mask to avoid performing attention on certain token indices by default a causal mask will be used to make sure the model can only look at previous inputs in order to predict the future head_mask torch tensor of shape encoder_layers encoder_attention_heads optional mask to nullify selected heads of the attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoder_head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the cross attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoder_outputs tuple tuple torch floattensor optional tuple consists of last_hidden_state hidden_states optional and attentions optional last_hidden_state of shape batch_size sequence_length hidden_size optional is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder past_key_values tuple tuple torch floattensor optional returned when use_cache true is passed or when config use_cache true tuple of tuple torch floattensor of length config n_layers with each tuple having 2 tensors of shape batch_size num_heads sequence_length embed_size_per_head and 2 additional tensors of shape batch_size num_heads encoder_sequence_length embed_size_per_head contains pre computed hidden states key and values in the self attention blocks and in the cross attention blocks that can be used see past_key_values input to speed up sequential decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple transformer encoder consisting of config encoder_layers self attention layers each layer is a timeseriestransformerencoderlayer args config timeseriestransformerconfig initialize weights and apply final processing args attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask head_mask torch tensor of shape encoder_layers encoder_attention_heads optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple expand attention_mask bsz seq_len bsz 1 tgt_seq_len src_seq_len check if head_mask has a correct number of layers specified if desired add layerdrop see https arxiv org abs 1909 11556 for description skip the layer transformer decoder consisting of config decoder_layers layers each layer is a timeseriestransformerdecoderlayer args config timeseriestransformerconfig initialize weights and apply final processing args attention_mask torch tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask encoder_hidden_states torch floattensor of shape batch_size encoder_sequence_length hidden_size optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder encoder_attention_mask torch longtensor of shape batch_size encoder_sequence_length optional mask to avoid performing cross attention on padding tokens indices of encoder input_ids mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask torch tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the cross attention modules in the decoder to avoid performing cross attention on hidden heads mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked past_key_values tuple tuple torch floattensor optional returned when use_cache true is passed or when config use_cache true tuple of tuple torch floattensor of length config n_layers with each tuple having 2 tensors of shape batch_size num_heads sequence_length embed_size_per_head and 2 additional tensors of shape batch_size num_heads encoder_sequence_length embed_size_per_head contains pre computed hidden states key and values in the self attention blocks and in the cross attention blocks that can be used see past_key_values input to speed up sequential decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple past_key_values_length expand encoder attention mask bsz seq_len bsz 1 tgt_seq_len src_seq_len decoder layers check if head_mask cross_attn_head_mask has a correct number of layers specified if desired add layerdrop see https arxiv org abs 1909 11556 for description add hidden states from the last decoder layer transformer encoder decoder and mask initializer initialize weights and apply final processing returns lagged subsequences of a given sequence returns a tensor of shape n s c i where s subsequences_length and i len indices containing lagged subsequences specifically lagged i j k sequence i indices k s j args sequence tensor the sequence from which lagged subsequences should be extracted shape n t c subsequences_length int length of the subsequences to be extracted shift int shift the lags by this amount back time feature target static features all features lagged features transformer inputs returns examples python from huggingface_hub import hf_hub_download import torch from transformers import timeseriestransformermodel file hf_hub_download repo_id hf internal testing tourism monthly batch filename train batch pt repo_type dataset batch torch load file model timeseriestransformermodel from_pretrained huggingface time series transformer tourism monthly during training one provides both past and future values as well as possible additional features outputs 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 outputs last_hidden_state if the user passed a tuple for encoder_outputs we wrap it in a basemodeloutput when return_dict true initialize weights of distribution_output and apply final processing returns examples python from huggingface_hub import hf_hub_download import torch from transformers import timeseriestransformerforprediction file hf_hub_download repo_id hf internal testing tourism monthly batch filename train batch pt repo_type dataset batch torch load file model timeseriestransformerforprediction from_pretrained huggingface time series transformer tourism monthly during training one provides both past and future values as well as possible additional features outputs 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 loss outputs loss loss backward during inference one only provides past values as well as possible additional features the model autoregressively generates future values outputs model generate 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 mean_prediction outputs sequences mean dim 1 outputs last_hidden_state loc is 3rd last and scale is 2nd last output greedily generate sequences of sample predictions from a model with a probability distribution head parameters past_values torch floattensor of shape batch_size sequence_length or batch_size sequence_length input_size past values of the time series that serve as context in order to predict the future the sequence size of this tensor must be larger than the context_length of the model since the model will use the larger size to construct lag features i e additional values from the past which are added in order to serve as extra context the sequence_length here is equal to config context_length max config lags_sequence which if no lags_sequence is configured is equal to config context_length 7 as by default the largest look back index in config lags_sequence is 7 the property _past_length returns the actual length of the past the past_values is what the transformer encoder gets as input with optional additional features such as static_categorical_features static_real_features past_time_features and lags optionally missing values need to be replaced with zeros and indicated via the past_observed_mask for multivariate time series the input_size 1 dimension is required and corresponds to the number of variates in the time series per time step past_time_features torch floattensor of shape batch_size sequence_length num_features required time features which the model internally will add to past_values these could be things like month of year day of the month etc encoded as vectors for instance as fourier features these could also be so called age features which basically help the model know at which point in life a time series is age features have small values for distant past time steps and increase monotonically the more we approach the current time step holiday features are also a good example of time features these features serve as the positional encodings of the inputs so contrary to a model like bert where the position encodings are learned from scratch internally as parameters of the model the time series transformer requires to provide additional time features the time series transformer only learns additional embeddings for static_categorical_features additional dynamic real covariates can be concatenated to this tensor with the caveat that these features must but known at prediction time the num_features here is equal to config num_time_features config num_dynamic_real_features future_time_features torch floattensor of shape batch_size prediction_length num_features required time features for the prediction window which the model internally will add to sampled predictions these could be things like month of year day of the month etc encoded as vectors for instance as fourier features these could also be so called age features which basically help the model know at which point in life a time series is age features have small values for distant past time steps and increase monotonically the more we approach the current time step holiday features are also a good example of time features these features serve as the positional encodings of the inputs so contrary to a model like bert where the position encodings are learned from scratch internally as parameters of the model the time series transformer requires to provide additional time features the time series transformer only learns additional embeddings for static_categorical_features additional dynamic real covariates can be concatenated to this tensor with the caveat that these features must but known at prediction time the num_features here is equal to config num_time_features config num_dynamic_real_features past_observed_mask torch booltensor of shape batch_size sequence_length or batch_size sequence_length input_size optional boolean mask to indicate which past_values were observed and which were missing mask values selected in 0 1 1 for values that are observed 0 for values that are missing i e nans that were replaced by zeros static_categorical_features torch longtensor of shape batch_size number of static categorical features optional optional static categorical features for which the model will learn an embedding which it will add to the values of the time series static categorical features are features which have the same value for all time steps static over time a typical example of a static categorical feature is a time series id static_real_features torch floattensor of shape batch_size number of static real features optional optional static real features which the model will add to the values of the time series static real features are features which have the same value for all time steps static over time a typical example of a static real feature is promotion information output_attentions bool optional whether or not to return the attentions tensors of all attention layers output_hidden_states bool optional whether or not to return the hidden states of all layers return sampletspredictionoutput where the outputs sequences tensor will have shape batch_size number of samples prediction_length or batch_size number of samples prediction_length input_size for multivariate predictions greedy decoding
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
from torch import nn

from ...activations import ACT2FN
from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask
from ...modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPastAndCrossAttentions,
    SampleTSPredictionOutput,
    Seq2SeqTSModelOutput,
    Seq2SeqTSPredictionOutput,
)
from ...modeling_utils import PreTrainedModel
from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
from .configuration_time_series_transformer import TimeSeriesTransformerConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "TimeSeriesTransformerConfig"


TIME_SERIES_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "huggingface/time-series-transformer-tourism-monthly",
    
]


class TimeSeriesFeatureEmbedder(nn.Module):
    

    def __init__(self, cardinalities: List[int], embedding_dims: List[int]) -> None:
        super().__init__()

        self.num_features = len(cardinalities)
        self.embedders = nn.ModuleList([nn.Embedding(c, d) for c, d in zip(cardinalities, embedding_dims)])

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        if self.num_features > 1:
            
            
            cat_feature_slices = torch.chunk(features, self.num_features, dim=-1)
        else:
            cat_feature_slices = [features]

        return torch.cat(
            [
                embed(cat_feature_slice.squeeze(-1))
                for embed, cat_feature_slice in zip(self.embedders, cat_feature_slices)
            ],
            dim=-1,
        )


class TimeSeriesStdScaler(nn.Module):
    

    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__()
        self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
        self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
        self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-5

    def forward(
        self, data: torch.Tensor, observed_indicator: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        
        denominator = observed_indicator.sum(self.dim, keepdim=self.keepdim)
        denominator = denominator.clamp_min(1.0)
        loc = (data * observed_indicator).sum(self.dim, keepdim=self.keepdim) / denominator

        variance = (((data - loc) * observed_indicator) ** 2).sum(self.dim, keepdim=self.keepdim) / denominator
        scale = torch.sqrt(variance + self.minimum_scale)
        return (data - loc) / scale, loc, scale


class TimeSeriesMeanScaler(nn.Module):
    

    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__()
        self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
        self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
        self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-10
        self.default_scale = config.default_scale if hasattr(config, "default_scale") else None

    def forward(
        self, data: torch.Tensor, observed_indicator: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        
        ts_sum = (data * observed_indicator).abs().sum(self.dim, keepdim=True)
        num_observed = observed_indicator.sum(self.dim, keepdim=True)

        scale = ts_sum / torch.clamp(num_observed, min=1)

        
        
        if self.default_scale is None:
            batch_sum = ts_sum.sum(dim=0)
            batch_observations = torch.clamp(num_observed.sum(0), min=1)
            default_scale = torch.squeeze(batch_sum / batch_observations)
        else:
            default_scale = self.default_scale * torch.ones_like(scale)

        
        scale = torch.where(num_observed > 0, scale, default_scale)

        
        scale = torch.clamp(scale, min=self.minimum_scale)
        scaled_data = data / scale

        if not self.keepdim:
            scale = scale.squeeze(dim=self.dim)

        return scaled_data, torch.zeros_like(scale), scale


class TimeSeriesNOPScaler(nn.Module):
    

    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__()
        self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
        self.keepdim = config.keepdim if hasattr(config, "keepdim") else True

    def forward(
        self, data: torch.Tensor, observed_indicator: torch.Tensor = None
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        
        scale = torch.ones_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
        loc = torch.zeros_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
        return data, loc, scale


def nll(input: torch.distributions.Distribution, target: torch.Tensor) -> torch.Tensor:
    
    return -input.log_prob(target)


def weighted_average(input_tensor: torch.Tensor, weights: Optional[torch.Tensor] = None, dim=None) -> torch.Tensor:
    
    if weights is not None:
        weighted_tensor = torch.where(weights != 0, input_tensor * weights, torch.zeros_like(input_tensor))
        sum_weights = torch.clamp(weights.sum(dim=dim) if dim else weights.sum(), min=1.0)
        return (weighted_tensor.sum(dim=dim) if dim else weighted_tensor.sum()) / sum_weights
    else:
        return input_tensor.mean(dim=dim)



class TimeSeriesSinusoidalPositionalEmbedding(nn.Embedding):
    

    def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None) -> None:
        super().__init__(num_positions, embedding_dim)
        self.weight = self._init_weight(self.weight)

    @staticmethod
    def _init_weight(out: nn.Parameter) -> nn.Parameter:
        
        n_pos, dim = out.shape
        position_enc = np.array(
            [[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)]
        )
        out.requires_grad = False  
        sentinel = dim // 2 if dim % 2 == 0 else (dim // 2) + 1
        out[:, 0:sentinel] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))
        out[:, sentinel:] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))
        out.detach_()
        return out

    @torch.no_grad()
    def forward(self, input_ids_shape: torch.Size, past_key_values_length: int = 0) -> torch.Tensor:
        
        bsz, seq_len = input_ids_shape[:2]
        positions = torch.arange(
            past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
        )
        return super().forward(positions)


class TimeSeriesValueEmbedding(nn.Module):
    def __init__(self, feature_size, d_model):
        super().__init__()
        self.value_projection = nn.Linear(in_features=feature_size, out_features=d_model, bias=False)

    def forward(self, x):
        return self.value_projection(x)



class TimeSeriesTransformerAttention(nn.Module):
    

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        is_decoder: bool = False,
        bias: bool = True,
        is_causal: bool = False,
        config: Optional[TimeSeriesTransformerConfig] = None,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        self.config = config

        if (self.head_dim * num_heads) != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
                f" and `num_heads`: {num_heads})."
            )
        self.scaling = self.head_dim**-0.5
        self.is_decoder = is_decoder
        self.is_causal = is_causal

        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_value_states: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        layer_head_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        

        
        
        is_cross_attention = key_value_states is not None

        bsz, tgt_len, _ = hidden_states.size()

        
        query_states = self.q_proj(hidden_states) * self.scaling
        
        
        
        
        if (
            is_cross_attention
            and past_key_value is not None
            and past_key_value[0].shape[2] == key_value_states.shape[1]
        ):
            
            key_states = past_key_value[0]
            value_states = past_key_value[1]
        elif is_cross_attention:
            
            key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
            value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
        elif past_key_value is not None:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)
        else:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

        if self.is_decoder:
            
            
            
            
            
            
            
            past_key_value = (key_states, value_states)

        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
        key_states = key_states.reshape(*proj_shape)
        value_states = value_states.reshape(*proj_shape)

        src_len = key_states.size(1)
        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
                f" {attn_weights.size()}"
            )

        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)

        if layer_head_mask is not None:
            if layer_head_mask.size() != (self.num_heads,):
                raise ValueError(
                    f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
                    f" {layer_head_mask.size()}"
                )
            attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if output_attentions:
            
            
            
            
            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
        else:
            attn_weights_reshaped = None

        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

        attn_output = torch.bmm(attn_probs, value_states)

        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
        attn_output = attn_output.transpose(1, 2)

        
        
        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)

        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights_reshaped, past_key_value



class TimeSeriesTransformerEncoderLayer(nn.Module):
    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__()
        self.embed_dim = config.d_model
        attn_type = "flash_attention_2" if getattr(config, "_flash_attn_2_enabled", False) else "default"

        self.self_attn = TIME_SERIES_TRANSFORMER_ATTENTION_CLASSES[attn_type](
            embed_dim=self.embed_dim,
            num_heads=config.encoder_attention_heads,
            dropout=config.attention_dropout,
            config=config,
        )
        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout
        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = nn.LayerNorm(self.embed_dim)

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        attention_mask: torch.FloatTensor,
        layer_head_mask: torch.FloatTensor,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]:
        
        residual = hidden_states
        hidden_states, attn_weights, _ = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            layer_head_mask=layer_head_mask,
            output_attentions=output_attentions,
        )
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        residual = hidden_states
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.final_layer_norm(hidden_states)

        if hidden_states.dtype == torch.float16 and (
            torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
        ):
            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attn_weights,)

        return outputs


TIME_SERIES_TRANSFORMER_ATTENTION_CLASSES = {"default": TimeSeriesTransformerAttention}



class TimeSeriesTransformerDecoderLayer(nn.Module):
    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__()
        self.embed_dim = config.d_model

        attn_type = "flash_attention_2" if getattr(config, "_flash_attn_2_enabled", False) else "default"
        self.self_attn = TIME_SERIES_TRANSFORMER_ATTENTION_CLASSES[attn_type](
            embed_dim=self.embed_dim,
            num_heads=config.decoder_attention_heads,
            dropout=config.attention_dropout,
            is_decoder=True,
            is_causal=True,
            config=config,
        )
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout

        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.encoder_attn = TIME_SERIES_TRANSFORMER_ATTENTION_CLASSES[attn_type](
            self.embed_dim,
            config.decoder_attention_heads,
            dropout=config.attention_dropout,
            is_decoder=True,
            config=config,
        )
        self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
        self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = nn.LayerNorm(self.embed_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        layer_head_mask: Optional[torch.Tensor] = None,
        cross_attn_layer_head_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = True,
    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        
        residual = hidden_states

        
        
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            past_key_value=self_attn_past_key_value,
            attention_mask=attention_mask,
            layer_head_mask=layer_head_mask,
            output_attentions=output_attentions,
        )
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        
        cross_attn_present_key_value = None
        cross_attn_weights = None
        if encoder_hidden_states is not None:
            residual = hidden_states

            
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
                hidden_states=hidden_states,
                key_value_states=encoder_hidden_states,
                attention_mask=encoder_attention_mask,
                layer_head_mask=cross_attn_layer_head_mask,
                past_key_value=cross_attn_past_key_value,
                output_attentions=output_attentions,
            )
            hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
            hidden_states = residual + hidden_states
            hidden_states = self.encoder_attn_layer_norm(hidden_states)

            
            present_key_value = present_key_value + cross_attn_present_key_value

        
        residual = hidden_states
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
        hidden_states = self.final_layer_norm(hidden_states)

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (self_attn_weights, cross_attn_weights)

        if use_cache:
            outputs += (present_key_value,)

        return outputs


class TimeSeriesTransformerPreTrainedModel(PreTrainedModel):
    config_class = TimeSeriesTransformerConfig
    base_model_prefix = "model"
    main_input_name = "past_values"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        std = self.config.init_std
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, TimeSeriesSinusoidalPositionalEmbedding):
            pass
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


TIME_SERIES_TRANSFORMER_START_DOCSTRING = r

TIME_SERIES_TRANSFORMER_INPUTS_DOCSTRING = r


class TimeSeriesTransformerEncoder(TimeSeriesTransformerPreTrainedModel):
    

    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__(config)

        self.dropout = config.dropout
        self.layerdrop = config.encoder_layerdrop
        if config.prediction_length is None:
            raise ValueError("The `prediction_length` config needs to be specified.")

        self.value_embedding = TimeSeriesValueEmbedding(feature_size=config.feature_size, d_model=config.d_model)
        self.embed_positions = TimeSeriesSinusoidalPositionalEmbedding(
            config.context_length + config.prediction_length, config.d_model
        )
        self.layers = nn.ModuleList([TimeSeriesTransformerEncoderLayer(config) for _ in range(config.encoder_layers)])
        self.layernorm_embedding = nn.LayerNorm(config.d_model)

        self.gradient_checkpointing = False
        
        self.post_init()

    def forward(
        self,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutput]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        hidden_states = self.value_embedding(inputs_embeds)
        embed_pos = self.embed_positions(inputs_embeds.size())

        hidden_states = self.layernorm_embedding(hidden_states + embed_pos)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

        
        if attention_mask is not None:
            
            attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        
        if head_mask is not None:
            if head_mask.size()[0] != (len(self.layers)):
                raise ValueError(
                    f"The head_mask should be specified for {len(self.layers)} layers, but it is for"
                    f" {head_mask.size()[0]}."
                )

        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)
            
            to_drop = False
            if self.training:
                dropout_probability = torch.rand([])
                if dropout_probability < self.layerdrop:  
                    to_drop = True

            if to_drop:
                layer_outputs = (None, None)
            else:
                if self.gradient_checkpointing and self.training:
                    layer_outputs = self._gradient_checkpointing_func(
                        encoder_layer.__call__,
                        hidden_states,
                        attention_mask,
                        (head_mask[idx] if head_mask is not None else None),
                        output_attentions,
                    )
                else:
                    layer_outputs = encoder_layer(
                        hidden_states,
                        attention_mask,
                        layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                        output_attentions=output_attentions,
                    )

                hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
        return BaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
        )


class TimeSeriesTransformerDecoder(TimeSeriesTransformerPreTrainedModel):
    

    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__(config)
        self.dropout = config.dropout
        self.layerdrop = config.decoder_layerdrop
        if config.prediction_length is None:
            raise ValueError("The `prediction_length` config needs to be specified.")

        self.value_embedding = TimeSeriesValueEmbedding(feature_size=config.feature_size, d_model=config.d_model)
        self.embed_positions = TimeSeriesSinusoidalPositionalEmbedding(
            config.context_length + config.prediction_length, config.d_model
        )
        self.layers = nn.ModuleList([TimeSeriesTransformerDecoderLayer(config) for _ in range(config.decoder_layers)])
        self.layernorm_embedding = nn.LayerNorm(config.d_model)

        self.gradient_checkpointing = False
        
        self.post_init()

    def forward(
        self,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        input_shape = inputs_embeds.size()[:-1]

        
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        attention_mask = _prepare_4d_causal_attention_mask(
            attention_mask, input_shape, inputs_embeds, past_key_values_length
        )

        
        if encoder_hidden_states is not None and encoder_attention_mask is not None:
            
            encoder_attention_mask = _prepare_4d_attention_mask(
                encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
            )

        hidden_states = self.value_embedding(inputs_embeds)
        embed_pos = self.embed_positions(inputs_embeds.size(), past_key_values_length=self.config.context_length)
        hidden_states = self.layernorm_embedding(hidden_states + embed_pos)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None
        next_decoder_cache = () if use_cache else None

        
        for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):
            if attn_mask is not None:
                if attn_mask.size()[0] != (len(self.layers)):
                    raise ValueError(
                        f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"
                        f" {head_mask.size()[0]}."
                    )

        for idx, decoder_layer in enumerate(self.layers):
            
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            if self.training:
                dropout_probability = torch.rand([])
                if dropout_probability < self.layerdrop:
                    continue

            past_key_value = past_key_values[idx] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    decoder_layer.__call__,
                    hidden_states,
                    attention_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    head_mask[idx] if head_mask is not None else None,
                    cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None,
                    None,
                    output_attentions,
                    use_cache,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                    cross_attn_layer_head_mask=(
                        cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None
                    ),
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                )
            hidden_states = layer_outputs[0]

            if use_cache:
                next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

                if encoder_hidden_states is not None:
                    all_cross_attentions += (layer_outputs[2],)

        
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None
        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
            cross_attentions=all_cross_attentions,
        )


@add_start_docstrings(
    "The bare Time Series Transformer Model outputting raw hidden-states without any specific head on top.",
    TIME_SERIES_TRANSFORMER_START_DOCSTRING,
)
class TimeSeriesTransformerModel(TimeSeriesTransformerPreTrainedModel):
    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__(config)

        if config.scaling == "mean" or config.scaling is True:
            self.scaler = TimeSeriesMeanScaler(config)
        elif config.scaling == "std":
            self.scaler = TimeSeriesStdScaler(config)
        else:
            self.scaler = TimeSeriesNOPScaler(config)

        if config.num_static_categorical_features > 0:
            self.embedder = TimeSeriesFeatureEmbedder(
                cardinalities=config.cardinality,
                embedding_dims=config.embedding_dimension,
            )

        
        self.encoder = TimeSeriesTransformerEncoder(config)
        self.decoder = TimeSeriesTransformerDecoder(config)

        
        self.post_init()

    @property
    def _past_length(self) -> int:
        return self.config.context_length + max(self.config.lags_sequence)

    def get_lagged_subsequences(
        self, sequence: torch.Tensor, subsequences_length: int, shift: int = 0
    ) -> torch.Tensor:
        
        sequence_length = sequence.shape[1]
        indices = [lag - shift for lag in self.config.lags_sequence]

        if max(indices) + subsequences_length > sequence_length:
            raise ValueError(
                f"lags cannot go further than history length, found lag {max(indices)} "
                f"while history length is only {sequence_length}"
            )

        lagged_values = []
        for lag_index in indices:
            begin_index = -lag_index - subsequences_length
            end_index = -lag_index if lag_index > 0 else None
            lagged_values.append(sequence[:, begin_index:end_index, ...])
        return torch.stack(lagged_values, dim=-1)

    def create_network_inputs(
        self,
        past_values: torch.Tensor,
        past_time_features: torch.Tensor,
        static_categorical_features: Optional[torch.Tensor] = None,
        static_real_features: Optional[torch.Tensor] = None,
        past_observed_mask: Optional[torch.Tensor] = None,
        future_values: Optional[torch.Tensor] = None,
        future_time_features: Optional[torch.Tensor] = None,
    ):
        
        time_feat = (
            torch.cat(
                (
                    past_time_features[:, self._past_length - self.config.context_length :, ...],
                    future_time_features,
                ),
                dim=1,
            )
            if future_values is not None
            else past_time_features[:, self._past_length - self.config.context_length :, ...]
        )

        
        if past_observed_mask is None:
            past_observed_mask = torch.ones_like(past_values)

        context = past_values[:, -self.config.context_length :]
        observed_context = past_observed_mask[:, -self.config.context_length :]
        _, loc, scale = self.scaler(context, observed_context)

        inputs = (
            (torch.cat((past_values, future_values), dim=1) - loc) / scale
            if future_values is not None
            else (past_values - loc) / scale
        )

        
        log_abs_loc = loc.abs().log1p() if self.config.input_size == 1 else loc.squeeze(1).abs().log1p()
        log_scale = scale.log() if self.config.input_size == 1 else scale.squeeze(1).log()
        static_feat = torch.cat((log_abs_loc, log_scale), dim=1)

        if static_real_features is not None:
            static_feat = torch.cat((static_real_features, static_feat), dim=1)
        if static_categorical_features is not None:
            embedded_cat = self.embedder(static_categorical_features)
            static_feat = torch.cat((embedded_cat, static_feat), dim=1)
        expanded_static_feat = static_feat.unsqueeze(1).expand(-1, time_feat.shape[1], -1)

        
        features = torch.cat((expanded_static_feat, time_feat), dim=-1)

        
        subsequences_length = (
            self.config.context_length + self.config.prediction_length
            if future_values is not None
            else self.config.context_length
        )
        lagged_sequence = self.get_lagged_subsequences(sequence=inputs, subsequences_length=subsequences_length)
        lags_shape = lagged_sequence.shape
        reshaped_lagged_sequence = lagged_sequence.reshape(lags_shape[0], lags_shape[1], -1)

        if reshaped_lagged_sequence.shape[1] != time_feat.shape[1]:
            raise ValueError(
                f"input length {reshaped_lagged_sequence.shape[1]} and time feature lengths {time_feat.shape[1]} does not match"
            )

        
        transformer_inputs = torch.cat((reshaped_lagged_sequence, features), dim=-1)

        return transformer_inputs, loc, scale, static_feat

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    @add_start_docstrings_to_model_forward(TIME_SERIES_TRANSFORMER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=Seq2SeqTSModelOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        past_values: torch.Tensor,
        past_time_features: torch.Tensor,
        past_observed_mask: torch.Tensor,
        static_categorical_features: Optional[torch.Tensor] = None,
        static_real_features: Optional[torch.Tensor] = None,
        future_values: Optional[torch.Tensor] = None,
        future_time_features: Optional[torch.Tensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[List[torch.FloatTensor]] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Seq2SeqTSModelOutput, Tuple]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_inputs, loc, scale, static_feat = self.create_network_inputs(
            past_values=past_values,
            past_time_features=past_time_features,
            past_observed_mask=past_observed_mask,
            static_categorical_features=static_categorical_features,
            static_real_features=static_real_features,
            future_values=future_values,
            future_time_features=future_time_features,
        )

        if encoder_outputs is None:
            enc_input = transformer_inputs[:, : self.config.context_length, ...]
            encoder_outputs = self.encoder(
                inputs_embeds=enc_input,
                head_mask=head_mask,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
        
        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
            encoder_outputs = BaseModelOutput(
                last_hidden_state=encoder_outputs[0],
                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
            )

        dec_input = transformer_inputs[:, self.config.context_length :, ...]
        decoder_outputs = self.decoder(
            inputs_embeds=dec_input,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_outputs[0],
            head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs + (loc, scale, static_feat)

        return Seq2SeqTSModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
            loc=loc,
            scale=scale,
            static_features=static_feat,
        )


@add_start_docstrings(
    "The Time Series Transformer Model with a distribution head on top for time-series forecasting.",
    TIME_SERIES_TRANSFORMER_START_DOCSTRING,
)
class TimeSeriesTransformerForPrediction(TimeSeriesTransformerPreTrainedModel):
    def __init__(self, config: TimeSeriesTransformerConfig):
        super().__init__(config)
        self.model = TimeSeriesTransformerModel(config)
        if config.distribution_output == "student_t":
            self.distribution_output = StudentTOutput(dim=config.input_size)
        elif config.distribution_output == "normal":
            self.distribution_output = NormalOutput(dim=config.input_size)
        elif config.distribution_output == "negative_binomial":
            self.distribution_output = NegativeBinomialOutput(dim=config.input_size)
        else:
            raise ValueError(f"Unknown distribution output {config.distribution_output}")

        self.parameter_projection = self.distribution_output.get_parameter_projection(self.model.config.d_model)
        self.target_shape = self.distribution_output.event_shape

        if config.loss == "nll":
            self.loss = nll
        else:
            raise ValueError(f"Unknown loss function {config.loss}")

        
        self.post_init()

    def output_params(self, dec_output):
        return self.parameter_projection(dec_output)

    def get_encoder(self):
        return self.model.get_encoder()

    def get_decoder(self):
        return self.model.get_decoder()

    @torch.jit.ignore
    def output_distribution(self, params, loc=None, scale=None, trailing_n=None) -> torch.distributions.Distribution:
        sliced_params = params
        if trailing_n is not None:
            sliced_params = [p[:, -trailing_n:] for p in params]
        return self.distribution_output.distribution(sliced_params, loc=loc, scale=scale)

    @add_start_docstrings_to_model_forward(TIME_SERIES_TRANSFORMER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=Seq2SeqTSModelOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        past_values: torch.Tensor,
        past_time_features: torch.Tensor,
        past_observed_mask: torch.Tensor,
        static_categorical_features: Optional[torch.Tensor] = None,
        static_real_features: Optional[torch.Tensor] = None,
        future_values: Optional[torch.Tensor] = None,
        future_time_features: Optional[torch.Tensor] = None,
        future_observed_mask: Optional[torch.Tensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[List[torch.FloatTensor]] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Seq2SeqTSModelOutput, Tuple]:
        r

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if future_values is not None:
            use_cache = False

        outputs = self.model(
            past_values=past_values,
            past_time_features=past_time_features,
            past_observed_mask=past_observed_mask,
            static_categorical_features=static_categorical_features,
            static_real_features=static_real_features,
            future_values=future_values,
            future_time_features=future_time_features,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            encoder_outputs=encoder_outputs,
            past_key_values=past_key_values,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            use_cache=use_cache,
            return_dict=return_dict,
        )

        prediction_loss = None
        params = None
        if future_values is not None:
            params = self.output_params(outputs[0])  
            
            distribution = self.output_distribution(params, loc=outputs[-3], scale=outputs[-2])

            loss = self.loss(distribution, future_values)

            if future_observed_mask is None:
                future_observed_mask = torch.ones_like(future_values)

            if len(self.target_shape) == 0:
                loss_weights = future_observed_mask
            else:
                loss_weights, _ = future_observed_mask.min(dim=-1, keepdim=False)

            prediction_loss = weighted_average(loss, weights=loss_weights)

        if not return_dict:
            outputs = ((params,) + outputs[1:]) if params is not None else outputs[1:]
            return ((prediction_loss,) + outputs) if prediction_loss is not None else outputs

        return Seq2SeqTSPredictionOutput(
            loss=prediction_loss,
            params=params,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
            loc=outputs.loc,
            scale=outputs.scale,
            static_features=outputs.static_features,
        )

    @torch.no_grad()
    def generate(
        self,
        past_values: torch.Tensor,
        past_time_features: torch.Tensor,
        future_time_features: torch.Tensor,
        past_observed_mask: Optional[torch.Tensor] = None,
        static_categorical_features: Optional[torch.Tensor] = None,
        static_real_features: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
    ) -> SampleTSPredictionOutput:
        r
        outputs = self(
            static_categorical_features=static_categorical_features,
            static_real_features=static_real_features,
            past_time_features=past_time_features,
            past_values=past_values,
            past_observed_mask=past_observed_mask,
            future_time_features=future_time_features,
            future_values=None,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            use_cache=True,
        )

        decoder = self.model.get_decoder()
        enc_last_hidden = outputs.encoder_last_hidden_state
        loc = outputs.loc
        scale = outputs.scale
        static_feat = outputs.static_features

        num_parallel_samples = self.config.num_parallel_samples
        repeated_loc = loc.repeat_interleave(repeats=num_parallel_samples, dim=0)
        repeated_scale = scale.repeat_interleave(repeats=num_parallel_samples, dim=0)

        repeated_past_values = (
            past_values.repeat_interleave(repeats=num_parallel_samples, dim=0) - repeated_loc
        ) / repeated_scale

        expanded_static_feat = static_feat.unsqueeze(1).expand(-1, future_time_features.shape[1], -1)
        features = torch.cat((expanded_static_feat, future_time_features), dim=-1)
        repeated_features = features.repeat_interleave(repeats=num_parallel_samples, dim=0)

        repeated_enc_last_hidden = enc_last_hidden.repeat_interleave(repeats=num_parallel_samples, dim=0)

        future_samples = []

        
        for k in range(self.config.prediction_length):
            lagged_sequence = self.model.get_lagged_subsequences(
                sequence=repeated_past_values,
                subsequences_length=1 + k,
                shift=1,
            )

            lags_shape = lagged_sequence.shape
            reshaped_lagged_sequence = lagged_sequence.reshape(lags_shape[0], lags_shape[1], -1)

            decoder_input = torch.cat((reshaped_lagged_sequence, repeated_features[:, : k + 1]), dim=-1)

            dec_output = decoder(inputs_embeds=decoder_input, encoder_hidden_states=repeated_enc_last_hidden)
            dec_last_hidden = dec_output.last_hidden_state

            params = self.parameter_projection(dec_last_hidden[:, -1:])
            distr = self.output_distribution(params, loc=repeated_loc, scale=repeated_scale)
            next_sample = distr.sample()

            repeated_past_values = torch.cat(
                (repeated_past_values, (next_sample - repeated_loc) / repeated_scale), dim=1
            )
            future_samples.append(next_sample)

        concat_future_samples = torch.cat(future_samples, dim=1)

        return SampleTSPredictionOutput(
            sequences=concat_future_samples.reshape(
                (-1, num_parallel_samples, self.config.prediction_length) + self.target_shape,
            )
        )