codingutf8 2023 ibm hugging face 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 patchtst model import math from dataclasses import dataclass from typing import optional tuple union import torch from torch import nn from activations import act2cls from modelingoutputs import basemodeloutput from modelingutils import pretrainedmodel from timeseriesutils import negativebinomialoutput normaloutput studenttoutput from utils import modeloutput addstartdocstrings logging from configurationpatchtst import patchtstconfig logger logging getloggername configfordoc patchtstconfig patchtstpretrainedmodelarchivelist ibmpatchtstetth1pretrain see all patchtst models at https huggingface comodels filterpatchtst copied from transformers models bart modelingbart bartattention with bartpatchtst class patchtstattentionnn 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 class patchtstbatchnormnn module def initself config patchtstconfig super init self batchnorm nn batchnorm1dconfig dmodel epsconfig normeps def forwardself inputs torch tensor output inputs transpose1 2 output batchsize dmodel sequencelength output self batchnormoutput return output transpose1 2 def randommasking inputs torch tensor maskratio float unmaskedchannelindices list none channelconsistentmasking bool false maskvalue int 0 if maskratio 0 or maskratio 1 raise valueerrorfmask ratio maskratio has to be between 0 and 1 batchsize numchannels sequencelength numfeatures inputs shape device inputs device lenkeep intsequencelength 1 maskratio if channelconsistentmasking noise torch randbatchsize 1 sequencelength devicedevice noise in 0 1 bs x 1 x l noise noise repeat1 numchannels 1 bs x numchannels x time else noise in 0 1 bs x numchannels x l noise torch randbatchsize numchannels sequencelength devicedevice mask bs x numchannels x numpatch mask torch onesbatchsize numchannels sequencelength devicedevice mask lenkeep 0 sort noise for each sample idsshuffle torch argsortnoise dim1 ascend small is keep large is remove idsrestore torch argsortidsshuffle dim1 idsrestore bs x numchannels x l mask torch gathermask dim1 indexidsrestore mask mask unsqueeze1 repeat1 1 1 numfeatures mask bs x numchannels x numpatches x patchlength if unmaskedchannelindices is not none mask unmaskedchannelindices 0 inputsmask inputs maskedfillmask bool maskvalue return inputsmask mask 0 def forecastmasking inputs torch tensor numforecastmaskpatches unionlist int unmaskedchannelindices list none maskvalue int 0 if isinstancenumforecastmaskpatches int numforecastmaskpatches numforecastmaskpatches forecastmaskratios 1 for in numforecastmaskpatches batchsize numchannels sequencelength numfeatures inputs shape mask torch zerosbatchsize numchannels sequencelength deviceinputs device tlist totallength 0 totalratio sumforecastmaskratios for patchlength ratio in zipnumforecastmaskpatches forecastmaskratios if patchlength 0 or patchlength sequencelength raise valueerror fnumforecastmaskpatches patchlength should be greater than 0 and less than total patches templen intbatchsize ratio totalratio tlist appendpatchlength ratio templen totallength templen tlist sortedtlist keylambda x x2 if totallength batchsize tlist02 tlist02 batchsize totallength elif totallength batchsize tlist12 tlist12 totallength batchsize batch1 0 for patchlen templen in tlist batch2 batch1 templen maskbatch1 batch2 patchlen 1 batch1 batch2 perm torch randpermmask shape0 mask maskperm mask mask unsqueeze1 repeat1 1 1 numfeatures mask bs x numchannels x numpatch x patchlen if unmaskedchannelindices is not none mask unmaskedchannelindices 0 inputsmask inputs maskedfillmask bool maskvalue return inputsmask mask 0 class patchtstpatchifynn module def initself config patchtstconfig super init self sequencelength config contextlength self patchlength config patchlength self patchstride config patchstride if self sequencelength self patchlength raise valueerror fsequence length self sequencelength has to be greater than the patch length self patchlength get the number of patches self numpatches maxself sequencelength self patchlength self patchlength self patchstride 1 newsequencelength self patchlength self patchstride self numpatches 1 self sequencestart self sequencelength newsequencelength def forwardself pastvalues torch tensor sequencelength pastvalues shape2 if sequencelength self sequencelength raise valueerror finput sequence length sequencelength doesn t match model configuration self sequencelength output bs x newsequencelength x numchannels output pastvalues self sequencestart output bs x numpatches x numinputchannels x patchlength output output unfolddimension2 sizeself patchlength stepself patchstride output bs x numinputchannels x numpatches x patchlength output output transpose2 3 contiguous return output class patchtstmaskingnn module def initself config patchtstconfig super init self randommaskratio config randommaskratio self channelconsistentmasking config channelconsistentmasking self masktype config masktype self numforecastmaskpatches config numforecastmaskpatches self unmaskedchannelindices config unmaskedchannelindices self maskvalue config maskvalue if self unmaskedchannelindices is not none self unmaskedchannelindices sortedself unmaskedchannelindices def forwardself patchinput torch tensor if self masktype random maskedinput mask randommasking inputspatchinput maskratioself randommaskratio unmaskedchannelindicesself unmaskedchannelindices channelconsistentmaskingself channelconsistentmasking maskvalueself maskvalue elif self masktype forecast maskedinput mask forecastmasking inputspatchinput numforecastmaskpatchesself numforecastmaskpatches unmaskedchannelindicesself unmaskedchannelindices maskvalueself maskvalue else raise valueerrorfinvalid mask type self masktype mask bs x numinputchannels x numpatch mask mask bool return maskedinput mask class patchtstencoderlayernn module def initself config patchtstconfig super init self channelattention config channelattention multihead attention self selfattn patchtstattention embeddimconfig dmodel numheadsconfig numattentionheads dropoutconfig attentiondropout add norm of the sublayer 1 self dropoutpath1 nn dropoutconfig pathdropout if config pathdropout 0 else nn identity if config normtype batchnorm self normsublayer1 patchtstbatchnormconfig elif config normtype layernorm self normsublayer1 nn layernormconfig dmodel epsconfig normeps else raise valueerrorfconfig normtype is not a supported norm layer type add norm of the sublayer 2 if self channelattention self dropoutpath2 nn dropoutconfig pathdropout if config pathdropout 0 else nn identity if config normtype batchnorm self normsublayer2 patchtstbatchnormconfig elif config normtype layernorm self normsublayer2 nn layernormconfig dmodel epsconfig normeps else raise valueerrorfconfig normtype is not a supported norm layer type positionwise feedforward self ff nn sequential nn linearconfig dmodel config ffndim biasconfig bias act2clsconfig activationfunction nn dropoutconfig ffdropout if config ffdropout 0 else nn identity nn linearconfig ffndim config dmodel biasconfig bias add norm of sublayer 3 self dropoutpath3 nn dropoutconfig pathdropout if config pathdropout 0 else nn identity if config normtype batchnorm self normsublayer3 patchtstbatchnormconfig elif config normtype layernorm self normsublayer3 nn layernormconfig dmodel epsconfig normeps else raise valueerrorfconfig normtype is not a supported norm layer type self prenorm config prenorm def forwardself hiddenstate torch tensor outputattentions optionalbool none batchsize numinputchannels sequencelength dmodel hiddenstate shape first sublayer attention across time hiddenstates bsnumchannels x sequencelength x dmodel hiddenstate hiddenstate viewbatchsize numinputchannels sequencelength dmodel if self prenorm norm and multihead attention and add residual connection attnoutput attnweights self selfattn hiddenstatesself normsublayer1hiddenstate outputattentionsoutputattentions add residual connection with residual dropout hiddenstate hiddenstate self dropoutpath1attnoutput else multihead attention and add residual connection and norm standard transformer from bert attnoutput attnweights self selfattn hiddenstateshiddenstate outputattentionsoutputattentions hiddenstates bsnumchannels x sequencelength x dmodel hiddenstate self normsublayer1hiddenstate self dropoutpath1attnoutput hiddenstate bs x numchannels x sequencelength x dmodel hiddenstate hiddenstate reshapebatchsize numinputchannels sequencelength dmodel second sublayer attention across variable at any given time if self channelattention hiddenstate bs x sequencelength x numchannels x dmodel hiddenstate hiddenstate transpose2 1 contiguous hiddenstate bssequencelength x numchannels x dmodel hiddenstate hiddenstate viewbatchsize sequencelength numinputchannels dmodel if self prenorm norm and multihead attention and add residual connection attnoutput channelattnweights self selfattn hiddenstatesself normsublayer2hiddenstate outputattentionsoutputattentions add residual connection with residual dropout hiddenstate hiddenstate self dropoutpath2attnoutput else multihead attention and add residual connection and norm attnoutput channelattnweights self selfattn hiddenstateshiddenstate outputattentionsoutputattentions hiddenstates bssequencelength x numchannels x dmodel hiddenstate self normsublayer2hiddenstate self dropoutpath2attnoutput reshape hidden state hiddenstate bs x sequencelength x numchannels x dmodel hiddenstate hiddenstate reshapebatchsize sequencelength numinputchannels dmodel hiddenstate bs x numchannels x sequencelength x dmodel hiddenstate hiddenstate transpose1 2 contiguous third sublayer mixing across hidden hiddenstate batchsizenumchannels x sequencelength x dmodel hiddenstate hiddenstate viewbatchsize numinputchannels sequencelength dmodel if self prenorm norm and positionwise feedforward and add residual connection add residual connection with residual dropout hiddenstate hiddenstate self dropoutpath3self ffself normsublayer3hiddenstate else positionwise feedforward and add residual connection and norm add residual connection with residual dropout hiddenstate self normsublayer3hiddenstate self dropoutpath3self ffhiddenstate bs x numchannels x sequencelength x dmodel hiddenstate hiddenstate reshapebatchsize numinputchannels sequencelength dmodel outputs hiddenstate if outputattentions outputs attnweights channelattnweights if self channelattention else attnweights return outputs class patchtstpretrainedmodelpretrainedmodel configclass patchtstconfig basemodelprefix model maininputname pastvalues supportsgradientcheckpointing false def initweightsself module if isinstancemodule patchtstpositionalencoding initialize clstoken if self config useclstoken nn init normalmodule clstoken std0 02 initialize positional encoding if self config positionalencodingtype random nn init normalmodule positionenc mean0 0 std0 1 elif isinstancemodule nn layernorm module bias data zero module weight data fill1 0 elif isinstancemodule patchtstbatchnorm module batchnorm bias data zero module batchnorm weight data fill1 0 elif isinstancemodule nn linear nn conv1d module weight data normalmean0 0 stdself config initstd if module bias is not none module bias data zero def setgradientcheckpointingself module valuefalse if isinstancemodule patchtstencoder module gradientcheckpointing value class patchtstembeddingnn module def initself config patchtstconfig super init self numinputchannels config numinputchannels self shareembedding config shareembedding input encoding projection of feature vectors onto a ddim vector space if self shareembedding self inputembedding nn linearconfig patchlength config dmodel else self inputembedding nn modulelist for in rangeconfig numinputchannels self inputembedding appendnn linearconfig patchlength config dmodel def forwardself patchinput torch tensor input encoding numinputchannels patchinput shape1 if numinputchannels self numinputchannels raise valueerror fthe defined number of input channels self numinputchannels in the config fhas to be the same as the number of channels in the batch input numinputchannels if self shareembedding embeddings self inputembeddingpatchinput x bs x numchannels x numpatches x dmodel else embeddings self inputembeddingipatchinput i for i in rangenuminputchannels embeddings torch stackembeddings dim1 return embeddings class patchtstpositionalencodingnn module def initself config patchtstconfig numpatches int super init self useclstoken config useclstoken self numinputchannels config numinputchannels if config useclstoken clstoken 1 x numinputchannels x 1 x dmodel self clstoken nn parametertorch zeros1 1 1 config dmodel numpatches 1 postional encoding numpatches x dmodel self positionenc self initpeconfig numpatches positional dropout self positionaldropout nn dropoutconfig positionaldropout if config positionaldropout 0 else nn identity staticmethod def initpeconfig patchtstconfig numpatches int nn parameter positional encoding if config positionalencodingtype random positionenc nn parametertorch randnnumpatches config dmodel requiresgradtrue elif config positionalencodingtype sincos positionenc torch zerosnumpatches config dmodel position torch arange0 numpatches unsqueeze1 divterm torch exptorch arange0 config dmodel 2 math log10000 0 config dmodel positionenc 0 2 torch sinposition divterm positionenc 1 2 torch cosposition divterm positionenc positionenc positionenc mean positionenc positionenc positionenc std 10 positionenc nn parameterpositionenc requiresgradfalse else raise valueerror fconfig positionalencodingtype is not a valid positional encoder available types are random and sincos return positionenc def forwardself patchinput torch tensor if self useclstoken patchinput bs x numchannels x numpatches x dmodel patchinput self positionaldropoutpatchinput self positionenc1 append cls token where clstoken 1 x numchannels x 1 x dmodel clstoken self clstoken self positionenc 1 get the same copy of clstoken for all the samples in batch bs x numchannels x 1 x dmodel clstokens clstoken expandpatchinput shape0 self numinputchannels 1 1 hiddenstate bs x numchannels x numpatches1 x dmodel hiddenstate torch catclstokens patchinput dim2 else hiddenstate bs x numchannels x numpatches x dmodel hiddenstate self positionaldropoutpatchinput self positionenc return hiddenstate class patchtstencoderpatchtstpretrainedmodel def initself config patchtstconfig numpatches int super initconfig self gradientcheckpointing false input embedding projection of feature vectors onto a ddim vector space self embedder patchtstembeddingconfig positional encoding self positionalencoder patchtstpositionalencodingconfig numpatches encoder self layers nn modulelistpatchtstencoderlayerconfig for i in rangeconfig numhiddenlayers initialize weights and apply final processing self postinit def forward self patchinput torch tensor outputhiddenstates optionalbool none outputattentions optionalbool none basemodeloutput outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates input embedding patchinput self embedderpatchinput positional encoding hiddenstate self positionalencoderpatchinput encoderstates if outputhiddenstates else none allattentions if outputattentions else none for encoderlayer in self layers if outputhiddenstates encoderstates encoderstates hiddenstate layeroutputs encoderlayerhiddenstatehiddenstate outputattentionsoutputattentions get hidden state hiddenstate shape is bs x numchannels x numpatches x dmodel or bs x numchannels x numpatches1 x dmodel if use clstoken hiddenstate layeroutputs0 append attention matrix at each layer if outputattentions allattentions allattentions layeroutputs1 return pastvalues hiddenstates return basemodeloutputlasthiddenstatehiddenstate hiddenstatesencoderstates attentionsallattentions patchtststartdocstring 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 patchtstconfig 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 base class for model s outputs with potential hidden states parameters lasthiddenstate torch floattensor of shape batchsize numchannels numpatches patchlength sequence of hiddenstates at the output of the last layer of the model hiddenstates tupletorch floattensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of torch floattensor one for the output of the embeddings if the model has an embedding layer one for the output of each layer of shape batchsize numchannels height width hiddenstates of the model at the output of each layer plus the optional initial embedding outputs mask torch floattensor of shape batchsize numchannels numpatches optional bool masked tensor indicating which patches are masked loc torch floattensor of shape batchsize 1 numchannels optional mean of the input data batchsize sequencelength numchannels over the sequencelength scale torch floattensor of shape batchsize 1 numchannels optional std of the input data batchsize sequencelength numchannels over the sequencelength patchinput torch floattensor of shape batchsize numchannels numpatches patchlength patched input to the transformer output type of patchtstforpretraining parameters loss optional returned when labels is provided torch floattensor of shape 1 mse loss predictionoutputs torch floattensor of shape batchsize sequencelength config vocabsize prediction outputs of the time series modeling heads hiddenstates tupletorch floattensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the initial embedding outputs attentions tupletorch floattensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of torch floattensor one for each layer of shape batchsize numheads sequencelength sequencelength attentions weights after the attention softmax used to compute the weighted average in the selfattention heads output type of patchtstforregression parameters loss optional returned when labels is provided torch floattensor of shape 1 mse loss regressionoutputs torch floattensor of shape batchsize numtargets regression outputs of the time series modeling heads hiddenstates tupletorch floattensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the initial embedding outputs attentions tupletorch floattensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of torch floattensor one for each layer of shape batchsize numheads sequencelength sequencelength attentions weights after the attention softmax used to compute the weighted average in the selfattention heads output type of patchtstforprediction parameters loss optional returned when labels is provided torch floattensor of shape 1 mse loss predictionoutputs torch floattensor of shape batchsize predictionlength 1 prediction outputs of the time series modeling heads hiddenstates tupletorch floattensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the initial embedding outputs attentions tupletorch floattensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of torch floattensor one for each layer of shape batchsize numheads sequencelength sequencelength attentions weights after the attention softmax used to compute the weighted average in the selfattention heads loc torch floattensor of shape batchsize 1 numchannels optional mean of the input data batchsize sequencelength numchannels over the sequencelength scale torch floattensor of shape batchsize 1 numchannels optional std of the input data batchsize sequencelength numchannels over the sequencelength output type of patchtstforclassification parameters loss optional returned when labels is provided torch floattensor of shape 1 total loss as the sum of the masked language modeling loss and the next sequence prediction classification loss predictionlogits torch floattensor of shape batchsize numtargets prediction scores of the patchtst modeling head scores before softmax hiddenstates tupletorch floattensor optional returned when outputhiddenstatestrue is passed or when config outputhiddenstatestrue tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batchsize sequencelength hiddensize hiddenstates of the model at the output of each layer plus the initial embedding outputs attentions tupletorch floattensor optional returned when outputattentionstrue is passed or when config outputattentionstrue tuple of torch floattensor one for each layer of shape batchsize numheads sequencelength sequencelength attentions weights after the attention softmax used to compute the weighted average in the selfattention heads base class for time series model s predictions outputs that contains the sampled values from the chosen distribution parameters sequences batchsize numsamples predictionlength numtargets sampled values from the chosen distribution copied from transformers models timeseriestransformer modelingtimeseriestransformer nll computes the negative log likelihood loss from input distribution with respect to target copied from transformers models timeseriestransformer modelingtimeseriestransformer weightedaverage 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 inputtensor torch floattensor input tensor of which the average must be computed weights torch floattensor optional weights tensor of the same shape as inputtensor dim int optional the dim along which to average inputtensor returns torch floattensor the tensor with values averaged along the specified dim copied from transformers models timeseriestransformer modelingtimeseriestransformer timeseriesstdscaler with timeseriestransformerpatchtst timeseriespatchtst 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 batchsize sequencelength numinputchannels input for batch norm calculation observedindicator torch booltensor of shape batchsize sequencelength numinputchannels calculating the scale on the observed indicator returns tuple of torch tensor of shapes batchsize sequencelength numinputchannels batchsize 1 numinputchannels batchsize 1 numinputchannels copied from transformers models timeseriestransformer modelingtimeseriestransformer timeseriesmeanscaler with timeseriestransformerpatchtst timeseriespatchtst 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 batchsize sequencelength numinputchannels input for batch norm calculation observedindicator torch booltensor of shape batchsize sequencelength numinputchannels calculating the scale on the observed indicator returns tuple of torch tensor of shapes batchsize sequencelength numinputchannels batchsize 1 numinputchannels batchsize 1 numinputchannels if defaultscale 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 minimumscale copied from transformers models timeseriestransformer modelingtimeseriestransformer timeseriesnopscaler with timeseriestransformerpatchtst timeseriespatchtst 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 batchsize sequencelength numinputchannels input for batch norm calculation returns tuple of torch tensor of shapes batchsize sequencelength numinputchannels batchsize 1 numinputchannels batchsize 1 numinputchannels parameters data torch tensor of shape batchsize sequencelength numinputchannels input for scaler calculation observedindicator torch booltensor of shape batchsize sequencelength numinputchannels calculating the scale on the observed indicator returns tuple of torch tensor of shapes batchsize sequencelength numinputchannels batchsize 1 numinputchannels batchsize 1 uminputchannels get numpatches information from patchtstpatchify initialize weights and apply final processing x tensor bs x sequencelength x numinputchannels patchedvalues bs x numinputchannels x numpatches x patchlength for pretrain pretraining head for mask modelling parameters embedding torch tensor of shape bs numchannels numpatches dmodel or bs numchannels numpatches1 dmodel if clstoken is set to true required embedding from the model returns torch tensor of shape bs numchannels numpatches dmodel or bs numchannels numpatches1 dmodel if clstoken is set to true initialize weights and apply final processing pastvalues bs x numchannels x numpatches x dmodel or bs x numchannels x numpatches1 x dmodel if use clstoken lasthiddenstate bs x numchannels x numpatches x patchlength or bs x numchannels x numpatches1 x patchlength if use clstoken calculate maskedloss parameters embedding torch tensor of shape bs numchannels numpatches dmodel or bs numchannels numpatches1 dmodel if clstoken is set to true required embedding from the model returns torch tensor of shape bs numtargets use the first output token pooledembedding bs x numchannels x dmodel pooledembedding bs x numchannels x dmodel pooledembedding bs x numchannels x dmodel pooledembedding bs x numchannels dmodel output bs x nclasses turn off masking initialize weights and apply final processing if each channel has its own head use linear head use distribution head all the channels share the same head use linear head use distribution head parameters embedding torch tensor of shape bs numchannels numpatches dmodel or bs numchannels numpatches1 dmodel if clstoken is set to true required embedding from the model returns torch tensor of shape bs forecastlen numchannels pooledembedding bs x numchannels x dmodel pooledembedding bs x numchannels x dmodel pooledembedding bs x numchannels x dmodel pooledembedding bs x numchannels x numpatches x dmodel pooledembedding bs x dmodel numpatches or bs x dmodel pooledembedding bs x forecastlen or tuple bs x forecastlen bs x forecastlen if using distribution head output bs x numchannels x forecastlen pooledembedding bs x numchannels x dmodel numpatches or bs x numchannels x dmodel output bs x numchannels x forecastlen or tuple bs x numchannels x forecastlen bs x numchannels x forecastlen if using distribution head output bs x forecastlen x numchannels bs x forecastlen x numchannels turn off masking initialize weights and apply final processing get model output get output head take average of the loss generate sequences of sample predictions from a model with a probability distribution head parameters pastvalues torch floattensor of shape batchsize sequencelength numinputchannels past values of the time series that serves as context in order to predict the future pastobservedmask torch booltensor of shape batchsize sequencelength numinputchannels 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 return samplepatchtstoutput where the outputs sequences tensor will have shape batchsize number of samples predictionlength 1 or batchsize number of samples predictionlength numinputchannels for multivariate predictions get number of samples get model output get distribution get samples list of bs x forecastlen x numchannels samples bs x numsamples x forecastlen x numchannels regression head parameters embedding torch tensor of shape bs numchannels numpatches dmodel or bs numchannels numpatches1 dmodel if clstoken is set to true required embedding from the model returns torch tensor of shape bs outputdim use the first output token pooledembedding bs x numchannels x dmodel pooledembedding bs x numchannels x dmodel pooledembedding bs x numchannels x dmodel flatten the input pooledembedding bs x numchannels dmodel projection output bs x outputdim or a tuple of this shape for distribution head apply sigmoid to bound the output if required turn off masking initialize weights and apply final processing get output head yhat is of shape bs x numtargets or tuple of this shape take average of the loss generate sequences of sample predictions from a model with a probability distribution head parameters pastvalues torch floattensor of shape batchsize sequencelength numinputchannels past values of the time series that serves as context in order to predict the future pastobservedmask torch booltensor of shape batchsize sequencelength numinputchannels 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 return samplepatchtstoutput where the outputs sequences tensor will have shape batchsize number of samples numtargets get number of samples get model output get distribution get samples list of bs x numtargets samples bs x numsamples x numtargets coding utf 8 2023 ibm hugging face 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 patchtst model see all patchtst models at https huggingface co models filter patchtst copied from transformers models bart modeling_bart bartattention with bart patchtst 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 compute batch normalization over the sequence length time dimension parameters inputs torch tensor of shape batch_size sequence_length d_model input for batch norm calculation returns torch tensor of shape batch_size sequence_length d_model output batch_size d_model sequence_length random_masking mask the input considering the control variables args inputs torch tensor of shape batch_size num_channels sequence_length num_features the input tensor to mask mask_ratio float masking ratio applied to mask the input data during random pretraining it is the number between 0 and 1 unmasked_channel_indices list optional indices of channels that will not be masked channel_consistent_masking bool optional defaults to false when true masking will be same across all channels of a timeseries otherwise masking positions will vary across channels mask_value int optional defaults to 0 define the value of masked patches for pretraining returns tuple torch tensor inputs_mask masked input same shape as input tensor and mask tensor of shape bs x c x n noise in 0 1 bs x 1 x l bs x num_channels x time noise in 0 1 bs x num_channels x l mask bs x num_channels x num_patch sort noise for each sample ascend small is keep large is remove ids_restore bs x num_channels x l mask bs x num_channels x num_patches x patch_length forecast masking that masks the last k patches where k is from the num_forecast_mask_patches if num_forecast_mask_patches is a list samples in the batch will be randomly masked by numbers defined in the list parameters inputs torch tensor input of shape bs num_channels num_patch patch_len num_forecast_mask_patches list number of patches to be masked at the end of each batch sample e g 4 or 3 5 unmasked_channel_indices list optional indices of channels that are not masked mask_value int optional defaults to 0 values in the masked patches will be filled by mask_value returns tuple torch tensor inputs_mask masked input same shape as inputs tensor and mask tensor of shape bs num_channels num_patch or bs tsg1 tsg2 num_channels num_patch mask bs x num_channels x num_patch x patch_len a class to patchify the time series sequence into different patches returns torch tensor of shape batch_size num_channels num_patches patch_length get the number of patches parameters past_values torch tensor of shape batch_size sequence_length num_channels required input for patchification returns torch tensor of shape batch_size num_channels num_patches patch_length output bs x new_sequence_length x num_channels output bs x num_patches x num_input_channels x patch_length output bs x num_input_channels x num_patches x patch_length class to perform random or forecast masking parameters config patchtstconfig model config returns x_mask torch tensor of shape batch_size num_channels num_patches patch_length masked patched input mask torch tensor of shape batch_size num_channels num_patches bool tensor indicating true on masked points parameters patch_input torch tensor of shape batch_size num_channels num_patches patch_length required patch input return masked_input torch tensor of shape batch_size num_channels num_patches patch_length masked patched input mask torch tensor of shape batch_size num_channels num_patches bool tensor indicating true on masked points mask bs x num_input_channels x num_patch patchtst encoder layer multi head attention add norm of the sublayer 1 add norm of the sublayer 2 position wise feed forward add norm of sublayer 3 parameters hidden_state torch tensor of shape batch_size num_channels sequence_length d_model required past values of the time series output_attentions bool optional whether or not to return the output attention of all layers return torch tensor of shape batch_size num_channels sequence_length d_model first sublayer attention across time hidden_states bs num_channels x sequence_length x d_model norm and multi head attention and add residual connection add residual connection with residual dropout multi head attention and add residual connection and norm standard transformer from bert hidden_states bs num_channels x sequence_length x d_model hidden_state bs x num_channels x sequence_length x d_model second sublayer attention across variable at any given time hidden_state bs x sequence_length x num_channels x d_model hidden_state bs sequence_length x num_channels x d_model norm and multi head attention and add residual connection add residual connection with residual dropout multi head attention and add residual connection and norm hidden_states bs sequence_length x num_channels x d_model reshape hidden state hidden_state bs x sequence_length x num_channels x d_model hidden_state bs x num_channels x sequence_length x d_model third sublayer mixing across hidden hidden_state batch_size num_channels x sequence_length x d_model norm and position wise feed forward and add residual connection add residual connection with residual dropout position wise feed forward and add residual connection and norm add residual connection with residual dropout bs x num_channels x sequence_length x d_model initialize weights initialize cls_token initialize positional encoding input encoding projection of feature vectors onto a d dim vector space parameters patch_input torch tensor of shape batch_size num_channels num_patches patch_length required patch input for embedding return torch tensor of shape batch_size num_channels num_patches d_model input encoding x bs x num_channels x num_patches x d_model class for positional encoding cls_token 1 x num_input_channels x 1 x d_model postional encoding num_patches x d_model positional dropout positional encoding patch_input bs x num_channels x num_patches x d_model append cls token where cls_token 1 x num_channels x 1 x d_model get the same copy of cls_token for all the samples in batch bs x num_channels x 1 x d_model hidden_state bs x num_channels x num_patches 1 x d_model hidden_state bs x num_channels x num_patches x d_model patchtst encoder input embedding projection of feature vectors onto a d dim vector space positional encoding encoder initialize weights and apply final processing parameters patch_input torch tensor of shape batch_size num_channels num_patches patch_length required past values of the time series output_hidden_states bool optional indicates if hidden states should be outputted output_attentions bool optional indicates if attentions should be outputted return basemodeloutput input embedding positional encoding get hidden state hidden_state shape is bs x num_channels x num_patches x d_model or bs x num_channels x num_patches 1 x d_model if use cls_token append attention matrix at each layer return past_values hidden_states 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 patchtstconfig 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 base class for model s outputs with potential hidden states parameters last_hidden_state torch floattensor of shape batch_size num_channels num_patches patch_length sequence of hidden states at the output of the last layer of the model hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings if the model has an embedding layer one for the output of each layer of shape batch_size num_channels height width hidden states of the model at the output of each layer plus the optional initial embedding outputs mask torch floattensor of shape batch_size num_channels num_patches optional bool masked tensor indicating which patches are masked loc torch floattensor of shape batch_size 1 num_channels optional mean of the input data batch_size sequence_length num_channels over the sequence_length scale torch floattensor of shape batch_size 1 num_channels optional std of the input data batch_size sequence_length num_channels over the sequence_length patch_input torch floattensor of shape batch_size num_channels num_patches patch_length patched input to the transformer output type of patchtstforpretraining parameters loss optional returned when labels is provided torch floattensor of shape 1 mse loss prediction_outputs torch floattensor of shape batch_size sequence_length config vocab_size prediction outputs of the time series modeling heads hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads output type of patchtstforregression parameters loss optional returned when labels is provided torch floattensor of shape 1 mse loss regression_outputs torch floattensor of shape batch_size num_targets regression outputs of the time series modeling heads hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads output type of patchtstforprediction parameters loss optional returned when labels is provided torch floattensor of shape 1 mse loss prediction_outputs torch floattensor of shape batch_size prediction_length 1 prediction outputs of the time series modeling heads hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads loc torch floattensor of shape batch_size 1 num_channels optional mean of the input data batch_size sequence_length num_channels over the sequence_length scale torch floattensor of shape batch_size 1 num_channels optional std of the input data batch_size sequence_length num_channels over the sequence_length output type of patchtstforclassification parameters loss optional returned when labels is provided torch floattensor of shape 1 total loss as the sum of the masked language modeling loss and the next sequence prediction classification loss prediction_logits torch floattensor of shape batch_size num_targets prediction scores of the patchtst modeling head scores before softmax hidden_states tuple torch floattensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of torch floattensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple torch floattensor optional returned when output_attentions true is passed or when config output_attentions true tuple of torch floattensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads base class for time series model s predictions outputs that contains the sampled values from the chosen distribution parameters sequences batch_size num_samples prediction_length num_targets sampled values from the chosen distribution copied from transformers models time_series_transformer modeling_time_series_transformer nll computes the negative log likelihood loss from input distribution with respect to target copied from transformers models time_series_transformer modeling_time_series_transformer weighted_average 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 time_series_transformer modeling_time_series_transformer timeseriesstdscaler with timeseriestransformer patchtst timeseries patchtst 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 copied from transformers models time_series_transformer modeling_time_series_transformer timeseriesmeanscaler with timeseriestransformer patchtst timeseries patchtst 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 copied from transformers models time_series_transformer modeling_time_series_transformer timeseriesnopscaler with timeseriestransformer patchtst timeseries patchtst 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 parameters data torch tensor of shape batch_size sequence_length num_input_channels input for scaler 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 um_input_channels get num_patches information from patchtstpatchify initialize weights and apply final processing parameters past_values torch tensor of shape bs sequence_length num_input_channels required input sequence to the model past_observed_mask torch booltensor of shape batch_size sequence_length num_input_channels 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 future_values torch booltensor of shape batch_size prediction_length num_input_channels optional future target values associated with the past_values output_hidden_states bool optional whether or not to return the hidden states of all layers output_attentions bool optional whether or not to return the output attention of all layers return_dict bool optional whether or not to return a modeloutput instead of a plain tuple returns patchtstmodeloutput or tuple of torch tensor if return_dict false or config return_dict false examples python from huggingface_hub import hf_hub_download import torch from transformers import patchtstmodel file hf_hub_download repo_id hf internal testing etth1 hourly batch filename train batch pt repo_type dataset batch torch load file model patchtstmodel from_pretrained namctin patchtst_etth1_pretrain during training one provides both past and future values outputs model past_values batch past_values future_values batch future_values last_hidden_state outputs last_hidden_state x tensor bs x sequence_length x num_input_channels patched_values bs x num_input_channels x num_patches x patch_length for pretrain pretraining head for mask modelling parameters embedding torch tensor of shape bs num_channels num_patches d_model or bs num_channels num_patches 1 d_model if cls_token is set to true required embedding from the model returns torch tensor of shape bs num_channels num_patches d_model or bs num_channels num_patches 1 d_model if cls_token is set to true bs x num_channels x num_patches x patch_length remove the first cls token initialize weights and apply final processing parameters past_values torch tensor of shape bs sequence_length num_input_channels required input sequence to the model past_observed_mask torch booltensor of shape batch_size sequence_length num_input_channels 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 output_hidden_states bool optional whether or not to return the hidden states of all layers output_attentions bool optional whether or not to return the output attention of all layers return_dict bool optional whether or not to return a modeloutput instead of a plain tuple returns patchtstforpretrainingoutput or tuple of torch tensor if return_dict false or config return_dict false examples python from huggingface_hub import hf_hub_download import torch from transformers import patchtstconfig patchtstforpretraining file hf_hub_download repo_id hf internal testing etth1 hourly batch filename train batch pt repo_type dataset batch torch load file config for random mask pretraining config patchtstconfig num_input_channels 7 context_length 512 patch_length 12 stride 12 mask_type random random_mask_ratio 0 4 use_cls_token true config for forecast mask pretraining config patchtstconfig num_input_channels 7 context_length 512 patch_length 12 stride 12 mask_type forecast num_forecast_mask_patches 5 use_cls_token true model patchtstforpretraining config during training one provides both past and future values outputs model past_values batch past_values loss outputs loss loss backward past_values bs x num_channels x num_patches x d_model or bs x num_channels x num_patches 1 x d_model if use cls_token last_hidden_state bs x num_channels x num_patches x patch_length or bs x num_channels x num_patches 1 x patch_length if use cls_token calculate masked_loss parameters embedding torch tensor of shape bs num_channels num_patches d_model or bs num_channels num_patches 1 d_model if cls_token is set to true required embedding from the model returns torch tensor of shape bs num_targets use the first output token pooled_embedding bs x num_channels x d_model pooled_embedding bs x num_channels x d_model pooled_embedding bs x num_channels x d_model pooled_embedding bs x num_channels d_model output bs x n_classes turn off masking initialize weights and apply final processing parameters past_values torch tensor of shape bs sequence_length num_input_channels required input sequence to the model target_values torch tensor optional labels associates with the past_values past_observed_mask torch booltensor of shape batch_size sequence_length num_input_channels 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 output_hidden_states bool optional whether or not to return the hidden states of all layers output_attentions bool optional whether or not to return the output attention of all layers return_dict bool optional whether or not to return a modeloutput instead of a plain tuple returns patchtstforclassificationoutput or tuple of torch tensor if return_dict false or config return_dict false examples python from transformers import patchtstconfig patchtstforclassification classification task with two input channel2 and 3 classes config patchtstconfig num_input_channels 2 num_targets 3 context_length 512 patch_length 12 stride 12 use_cls_token true model patchtstforclassification config config during inference one only provides past values past_values torch randn 20 512 2 outputs model past_values past_values labels outputs prediction_logits if each channel has its own head use linear head use distribution head all the channels share the same head use linear head use distribution head parameters embedding torch tensor of shape bs num_channels num_patches d_model or bs num_channels num_patches 1 d_model if cls_token is set to true required embedding from the model returns torch tensor of shape bs forecast_len num_channels pooled_embedding bs x num_channels x d_model pooled_embedding bs x num_channels x d_model pooled_embedding bs x num_channels x d_model pooled_embedding bs x num_channels x num_patches x d_model pooled_embedding bs x d_model num_patches or bs x d_model pooled_embedding bs x forecast_len or tuple bs x forecast_len bs x forecast_len if using distribution head output bs x num_channels x forecast_len pooled_embedding bs x num_channels x d_model num_patches or bs x num_channels x d_model output bs x num_channels x forecast_len or tuple bs x num_channels x forecast_len bs x num_channels x forecast_len if using distribution head output bs x forecast_len x num_channels bs x forecast_len x num_channels bs x forecast_len x num_channels turn off masking initialize weights and apply final processing parameters past_values torch tensor of shape bs sequence_length num_input_channels required input sequence to the model past_observed_mask torch booltensor of shape batch_size sequence_length num_input_channels 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 future_values torch tensor of shape bs forecast_len num_input_channels optional future target values associated with the past_values output_hidden_states bool optional whether or not to return the hidden states of all layers output_attentions bool optional whether or not to return the output attention of all layers return_dict bool optional whether or not to return a modeloutput instead of a plain tuple returns patchtstforpredictionoutput or tuple of torch tensor if return_dict false or config return_dict false examples python from huggingface_hub import hf_hub_download import torch from transformers import patchtstconfig patchtstforprediction file hf_hub_download repo_id hf internal testing etth1 hourly batch filename train batch pt repo_type dataset batch torch load file prediction task with 7 input channels and prediction length is 96 model patchtstforprediction from_pretrained namctin patchtst_etth1_forecast during training one provides both past and future values outputs model past_values batch past_values future_values batch future_values loss outputs loss loss backward during inference one only provides past values the model outputs future values outputs model past_values batch past_values prediction_outputs outputs prediction_outputs get model output get output head take average of the loss generate sequences of sample predictions from a model with a probability distribution head parameters past_values torch floattensor of shape batch_size sequence_length num_input_channels past values of the time series that serves as context in order to predict the future past_observed_mask torch booltensor of shape batch_size sequence_length num_input_channels 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 return samplepatchtstoutput where the outputs sequences tensor will have shape batch_size number of samples prediction_length 1 or batch_size number of samples prediction_length num_input_channels for multivariate predictions get number of samples get model output get distribution get samples list of bs x forecast_len x num_channels samples bs x num_samples x forecast_len x num_channels regression head parameters embedding torch tensor of shape bs num_channels num_patches d_model or bs num_channels num_patches 1 d_model if cls_token is set to true required embedding from the model returns torch tensor of shape bs output_dim use the first output token pooled_embedding bs x num_channels x d_model pooled_embedding bs x num_channels x d_model pooled_embedding bs x num_channels x d_model flatten the input pooled_embedding bs x num_channels d_model projection output bs x output_dim or a tuple of this shape for distribution head apply sigmoid to bound the output if required linear head turn off masking initialize weights and apply final processing parameters past_values torch tensor of shape bs sequence_length num_input_channels required input sequence to the model target_values torch tensor of shape bs num_input_channels target values associates with the past_values past_observed_mask torch booltensor of shape batch_size sequence_length num_input_channels 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 output_hidden_states bool optional whether or not to return the hidden states of all layers output_attentions bool optional whether or not to return the output attention of all layers return_dict bool optional whether or not to return a modeloutput instead of a plain tuple returns patchtstforregressionoutput or tuple of torch tensor if return_dict false or config return_dict false examples python from transformers import patchtstconfig patchtstforregression regression task with 6 input channels and regress 2 targets model patchtstforregression from_pretrained namctin patchtst_etth1_regression during inference one only provides past values the model outputs future values past_values torch randn 20 512 6 outputs model past_values past_values regression_outputs outputs regression_outputs get output head y_hat is of shape bs x num_targets or tuple of this shape take average of the loss generate sequences of sample predictions from a model with a probability distribution head parameters past_values torch floattensor of shape batch_size sequence_length num_input_channels past values of the time series that serves as context in order to predict the future past_observed_mask torch booltensor of shape batch_size sequence_length num_input_channels 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 return samplepatchtstoutput where the outputs sequences tensor will have shape batch_size number of samples num_targets get number of samples get model output get distribution get samples list of bs x num_targets samples bs x num_samples x num_targets
import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
from torch import nn

from ...activations import ACT2CLS
from ...modeling_outputs import BaseModelOutput
from ...modeling_utils import PreTrainedModel
from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput
from ...utils import ModelOutput, add_start_docstrings, logging
from .configuration_patchtst import PatchTSTConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "PatchTSTConfig"

PATCHTST_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "ibm/patchtst-etth1-pretrain",
    
]



class PatchTSTAttention(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[PatchTSTConfig] = 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 PatchTSTBatchNorm(nn.Module):
    

    def __init__(self, config: PatchTSTConfig):
        super().__init__()
        self.batchnorm = nn.BatchNorm1d(config.d_model, eps=config.norm_eps)

    def forward(self, inputs: torch.Tensor):
        
        output = inputs.transpose(1, 2)  
        output = self.batchnorm(output)
        return output.transpose(1, 2)


def random_masking(
    inputs: torch.Tensor,
    mask_ratio: float,
    unmasked_channel_indices: list = None,
    channel_consistent_masking: bool = False,
    mask_value: int = 0,
):
    
    if mask_ratio < 0 or mask_ratio >= 1:
        raise ValueError(f"Mask ratio {mask_ratio} has to be between 0 and 1.")

    batch_size, num_channels, sequence_length, num_features = inputs.shape
    device = inputs.device

    len_keep = int(sequence_length * (1 - mask_ratio))

    if channel_consistent_masking:
        noise = torch.rand(batch_size, 1, sequence_length, device=device)  
        noise = noise.repeat(1, num_channels, 1)  
    else:
        
        noise = torch.rand(batch_size, num_channels, sequence_length, device=device)

    
    mask = torch.ones(batch_size, num_channels, sequence_length, device=device)
    mask[:, :, :len_keep] = 0

    
    ids_shuffle = torch.argsort(noise, dim=-1)  
    ids_restore = torch.argsort(ids_shuffle, dim=-1)  

    mask = torch.gather(mask, dim=-1, index=ids_restore)
    mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features)  
    if unmasked_channel_indices is not None:
        mask[:, unmasked_channel_indices, :, :] = 0

    inputs_mask = inputs.masked_fill(mask.bool(), mask_value)
    return inputs_mask, mask[..., 0]


def forecast_masking(
    inputs: torch.Tensor,
    num_forecast_mask_patches: Union[list, int],
    unmasked_channel_indices: list = None,
    mask_value: int = 0,
):
    

    if isinstance(num_forecast_mask_patches, int):
        num_forecast_mask_patches = [num_forecast_mask_patches]
    forecast_mask_ratios = [1 for _ in num_forecast_mask_patches]

    batch_size, num_channels, sequence_length, num_features = inputs.shape
    mask = torch.zeros(batch_size, num_channels, sequence_length, device=inputs.device)

    t_list = []
    total_length = 0
    total_ratio = sum(forecast_mask_ratios)

    for patch_length, ratio in zip(num_forecast_mask_patches, forecast_mask_ratios):
        if patch_length <= 0 or patch_length >= sequence_length:
            raise ValueError(
                f"num_forecast_mask_patches {patch_length} should be greater than 0 and less than total patches."
            )
        temp_len = int(batch_size * ratio / total_ratio)
        t_list.append([patch_length, ratio, temp_len])
        total_length += temp_len

    t_list = sorted(t_list, key=lambda x: x[2])

    if total_length < batch_size:
        t_list[0][2] = t_list[0][2] + (batch_size - total_length)
    elif total_length > batch_size:
        t_list[-1][2] = t_list[-1][2] + (total_length - batch_size)

    batch1 = 0
    for patch_len, _, temp_len in t_list:
        batch2 = batch1 + temp_len
        mask[batch1:batch2, :, -patch_len:] = 1
        batch1 = batch2

    perm = torch.randperm(mask.shape[0])
    mask = mask[perm]

    mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features)  
    if unmasked_channel_indices is not None:
        mask[:, unmasked_channel_indices, :, :] = 0

    inputs_mask = inputs.masked_fill(mask.bool(), mask_value)
    return inputs_mask, mask[..., 0]


class PatchTSTPatchify(nn.Module):
    

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

        self.sequence_length = config.context_length
        self.patch_length = config.patch_length
        self.patch_stride = config.patch_stride

        if self.sequence_length <= self.patch_length:
            raise ValueError(
                f"Sequence length ({self.sequence_length}) has to be greater than the patch length ({self.patch_length})"
            )

        
        self.num_patches = (max(self.sequence_length, self.patch_length) - self.patch_length) // self.patch_stride + 1
        new_sequence_length = self.patch_length + self.patch_stride * (self.num_patches - 1)
        self.sequence_start = self.sequence_length - new_sequence_length

    def forward(self, past_values: torch.Tensor):
        
        sequence_length = past_values.shape[-2]
        if sequence_length != self.sequence_length:
            raise ValueError(
                f"Input sequence length ({sequence_length}) doesn't match model configuration ({self.sequence_length})."
            )
        
        output = past_values[:, self.sequence_start :, :]
        
        output = output.unfold(dimension=-2, size=self.patch_length, step=self.patch_stride)
        
        output = output.transpose(-2, -3).contiguous()
        return output


class PatchTSTMasking(nn.Module):
    

    def __init__(self, config: PatchTSTConfig):
        super().__init__()
        self.random_mask_ratio = config.random_mask_ratio
        self.channel_consistent_masking = config.channel_consistent_masking
        self.mask_type = config.mask_type
        self.num_forecast_mask_patches = config.num_forecast_mask_patches
        self.unmasked_channel_indices = config.unmasked_channel_indices
        self.mask_value = config.mask_value
        if self.unmasked_channel_indices is not None:
            self.unmasked_channel_indices = sorted(self.unmasked_channel_indices)

    def forward(self, patch_input: torch.Tensor):
        
        if self.mask_type == "random":
            masked_input, mask = random_masking(
                inputs=patch_input,
                mask_ratio=self.random_mask_ratio,
                unmasked_channel_indices=self.unmasked_channel_indices,
                channel_consistent_masking=self.channel_consistent_masking,
                mask_value=self.mask_value,
            )
        elif self.mask_type == "forecast":
            masked_input, mask = forecast_masking(
                inputs=patch_input,
                num_forecast_mask_patches=self.num_forecast_mask_patches,
                unmasked_channel_indices=self.unmasked_channel_indices,
                mask_value=self.mask_value,
            )
        else:
            raise ValueError(f"Invalid mask type {self.mask_type}.")

        
        mask = mask.bool()
        return masked_input, mask


class PatchTSTEncoderLayer(nn.Module):
    

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

        self.channel_attention = config.channel_attention
        
        self.self_attn = PatchTSTAttention(
            embed_dim=config.d_model,
            num_heads=config.num_attention_heads,
            dropout=config.attention_dropout,
        )

        
        self.dropout_path1 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()
        if config.norm_type == "batchnorm":
            self.norm_sublayer1 = PatchTSTBatchNorm(config)
        elif config.norm_type == "layernorm":
            self.norm_sublayer1 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
        else:
            raise ValueError(f"{config.norm_type} is not a supported norm layer type.")

        
        if self.channel_attention:
            self.dropout_path2 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()
            if config.norm_type == "batchnorm":
                self.norm_sublayer2 = PatchTSTBatchNorm(config)
            elif config.norm_type == "layernorm":
                self.norm_sublayer2 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
            else:
                raise ValueError(f"{config.norm_type} is not a supported norm layer type.")

        
        self.ff = nn.Sequential(
            nn.Linear(config.d_model, config.ffn_dim, bias=config.bias),
            ACT2CLS[config.activation_function](),
            nn.Dropout(config.ff_dropout) if config.ff_dropout > 0 else nn.Identity(),
            nn.Linear(config.ffn_dim, config.d_model, bias=config.bias),
        )

        
        self.dropout_path3 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()
        if config.norm_type == "batchnorm":
            self.norm_sublayer3 = PatchTSTBatchNorm(config)
        elif config.norm_type == "layernorm":
            self.norm_sublayer3 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
        else:
            raise ValueError(f"{config.norm_type} is not a supported norm layer type.")

        self.pre_norm = config.pre_norm

    def forward(self, hidden_state: torch.Tensor, output_attentions: Optional[bool] = None):
        
        batch_size, num_input_channels, sequence_length, d_model = hidden_state.shape

        
        
        hidden_state = hidden_state.view(batch_size * num_input_channels, sequence_length, d_model)

        if self.pre_norm:
            
            attn_output, attn_weights, _ = self.self_attn(
                hidden_states=self.norm_sublayer1(hidden_state), output_attentions=output_attentions
            )
            
            hidden_state = hidden_state + self.dropout_path1(attn_output)
        else:
            
            attn_output, attn_weights, _ = self.self_attn(
                hidden_states=hidden_state, output_attentions=output_attentions
            )
            
            hidden_state = self.norm_sublayer1(hidden_state + self.dropout_path1(attn_output))

        
        hidden_state = hidden_state.reshape(batch_size, num_input_channels, sequence_length, d_model)

        
        if self.channel_attention:
            
            hidden_state = hidden_state.transpose(2, 1).contiguous()
            
            hidden_state = hidden_state.view(batch_size * sequence_length, num_input_channels, d_model)
            if self.pre_norm:
                
                attn_output, channel_attn_weights, _ = self.self_attn(
                    hidden_states=self.norm_sublayer2(hidden_state), output_attentions=output_attentions
                )
                
                hidden_state = hidden_state + self.dropout_path2(attn_output)
            else:
                
                attn_output, channel_attn_weights, _ = self.self_attn(
                    hidden_states=hidden_state, output_attentions=output_attentions
                )
                
                hidden_state = self.norm_sublayer2(hidden_state + self.dropout_path2(attn_output))

            
            
            hidden_state = hidden_state.reshape(batch_size, sequence_length, num_input_channels, d_model)
            
            hidden_state = hidden_state.transpose(1, 2).contiguous()

        
        
        hidden_state = hidden_state.view(batch_size * num_input_channels, sequence_length, d_model)
        if self.pre_norm:
            
            
            hidden_state = hidden_state + self.dropout_path3(self.ff(self.norm_sublayer3(hidden_state)))
        else:
            
            
            hidden_state = self.norm_sublayer3(hidden_state + self.dropout_path3(self.ff(hidden_state)))

        
        hidden_state = hidden_state.reshape(batch_size, num_input_channels, sequence_length, d_model)

        outputs = (hidden_state,)
        if output_attentions:
            outputs += (attn_weights, channel_attn_weights) if self.channel_attention else (attn_weights,)

        return outputs


class PatchTSTPreTrainedModel(PreTrainedModel):
    config_class = PatchTSTConfig
    base_model_prefix = "model"
    main_input_name = "past_values"
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        
        if isinstance(module, PatchTSTPositionalEncoding):
            
            if self.config.use_cls_token:
                nn.init.normal_(module.cls_token, std=0.02)
            
            if self.config.positional_encoding_type == "random":
                nn.init.normal_(module.position_enc, mean=0.0, std=0.1)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        elif isinstance(module, PatchTSTBatchNorm):
            module.batchnorm.bias.data.zero_()
            module.batchnorm.weight.data.fill_(1.0)
        elif isinstance(module, (nn.Linear, nn.Conv1d)):
            module.weight.data.normal_(mean=0.0, std=self.config.init_std)
            if module.bias is not None:
                module.bias.data.zero_()

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, (PatchTSTEncoder)):
            module.gradient_checkpointing = value


class PatchTSTEmbedding(nn.Module):
    def __init__(self, config: PatchTSTConfig):
        super().__init__()
        self.num_input_channels = config.num_input_channels
        self.share_embedding = config.share_embedding
        
        if self.share_embedding:
            self.input_embedding = nn.Linear(config.patch_length, config.d_model)
        else:
            self.input_embedding = nn.ModuleList()
            for _ in range(config.num_input_channels):
                self.input_embedding.append(nn.Linear(config.patch_length, config.d_model))

    def forward(self, patch_input: torch.Tensor):
        
        
        num_input_channels = patch_input.shape[1]
        if num_input_channels != self.num_input_channels:
            raise ValueError(
                f"The defined number of input channels ({self.num_input_channels}) in the config "
                f"has to be the same as the number of channels in the batch input ({num_input_channels})"
            )
        if self.share_embedding:
            embeddings = self.input_embedding(patch_input)  
        else:
            embeddings = [self.input_embedding[i](patch_input[:, i, :, :]) for i in range(num_input_channels)]
            embeddings = torch.stack(embeddings, dim=1)
        return embeddings


class PatchTSTPositionalEncoding(nn.Module):
    

    def __init__(self, config: PatchTSTConfig, num_patches: int):
        super().__init__()
        self.use_cls_token = config.use_cls_token
        self.num_input_channels = config.num_input_channels
        if config.use_cls_token:
            
            self.cls_token = nn.Parameter(torch.zeros(1, 1, 1, config.d_model))
            num_patches += 1
        
        self.position_enc = self._init_pe(config, num_patches)
        
        self.positional_dropout = (
            nn.Dropout(config.positional_dropout) if config.positional_dropout > 0 else nn.Identity()
        )

    @staticmethod
    def _init_pe(config: PatchTSTConfig, num_patches: int) -> nn.Parameter:
        
        if config.positional_encoding_type == "random":
            position_enc = nn.Parameter(torch.randn(num_patches, config.d_model), requires_grad=True)
        elif config.positional_encoding_type == "sincos":
            position_enc = torch.zeros(num_patches, config.d_model)
            position = torch.arange(0, num_patches).unsqueeze(1)
            div_term = torch.exp(torch.arange(0, config.d_model, 2) * -(math.log(10000.0) / config.d_model))
            position_enc[:, 0::2] = torch.sin(position * div_term)
            position_enc[:, 1::2] = torch.cos(position * div_term)
            position_enc = position_enc - position_enc.mean()
            position_enc = position_enc / (position_enc.std() * 10)
            position_enc = nn.Parameter(position_enc, requires_grad=False)
        else:
            raise ValueError(
                f"{config.positional_encoding_type} is not a valid positional encoder. Available types are 'random' and 'sincos'."
            )
        return position_enc

    def forward(self, patch_input: torch.Tensor):
        if self.use_cls_token:
            
            patch_input = self.positional_dropout(patch_input + self.position_enc[1:, :])
            
            cls_token = self.cls_token + self.position_enc[:1, :]
            
            cls_tokens = cls_token.expand(patch_input.shape[0], self.num_input_channels, -1, -1)
            
            hidden_state = torch.cat((cls_tokens, patch_input), dim=2)
        else:
            
            hidden_state = self.positional_dropout(patch_input + self.position_enc)
        return hidden_state


class PatchTSTEncoder(PatchTSTPreTrainedModel):
    

    def __init__(self, config: PatchTSTConfig, num_patches: int):
        super().__init__(config)
        self.gradient_checkpointing = False

        
        self.embedder = PatchTSTEmbedding(config)
        
        self.positional_encoder = PatchTSTPositionalEncoding(config, num_patches)
        
        self.layers = nn.ModuleList([PatchTSTEncoderLayer(config) for i in range(config.num_hidden_layers)])

        
        self.post_init()

    def forward(
        self,
        patch_input: torch.Tensor,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
    ) -> BaseModelOutput:
        
        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
        )

        
        patch_input = self.embedder(patch_input)
        
        hidden_state = self.positional_encoder(patch_input)

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None
        for encoder_layer in self.layers:
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_state,)

            layer_outputs = encoder_layer(hidden_state=hidden_state, output_attentions=output_attentions)
            
            
            hidden_state = layer_outputs[0]
            
            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)
        
        return BaseModelOutput(last_hidden_state=hidden_state, hidden_states=encoder_states, attentions=all_attentions)


PATCHTST_START_DOCSTRING = r


@dataclass
class PatchTSTModelOutput(ModelOutput):
    

    last_hidden_state: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    mask: torch.FloatTensor = None
    loc: torch.FloatTensor = None
    scale: torch.FloatTensor = None
    patch_input: torch.FloatTensor = None


@dataclass
class PatchTSTForPretrainingOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    prediction_output: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class PatchTSTForRegressionOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    regression_outputs: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class PatchTSTForPredictionOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    prediction_outputs: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    loc: torch.FloatTensor = None
    scale: torch.FloatTensor = None


@dataclass
class PatchTSTForClassificationOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    prediction_logits: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class SamplePatchTSTOutput(ModelOutput):
    

    sequences: torch.FloatTensor = None



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 PatchTSTStdScaler(nn.Module):
    

    def __init__(self, config: PatchTSTConfig):
        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 PatchTSTMeanScaler(nn.Module):
    

    def __init__(self, config: PatchTSTConfig):
        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 PatchTSTNOPScaler(nn.Module):
    

    def __init__(self, config: PatchTSTConfig):
        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


class PatchTSTScaler(nn.Module):
    def __init__(self, config: PatchTSTConfig):
        super().__init__()
        if config.scaling == "mean" or config.scaling is True:
            self.scaler = PatchTSTMeanScaler(config)
        elif config.scaling == "std":
            self.scaler = PatchTSTStdScaler(config)
        else:
            self.scaler = PatchTSTNOPScaler(config)

    def forward(
        self, data: torch.Tensor, observed_indicator: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        
        data, loc, scale = self.scaler(data, observed_indicator)
        return data, loc, scale


@add_start_docstrings(
    "The bare PatchTST Model outputting raw hidden-states without any specific head.",
    PATCHTST_START_DOCSTRING,
)
class PatchTSTModel(PatchTSTPreTrainedModel):
    def __init__(self, config: PatchTSTConfig):
        super().__init__(config)

        self.scaler = PatchTSTScaler(config)
        self.patchifier = PatchTSTPatchify(config)
        self.do_mask_input = config.do_mask_input
        
        num_patches = self.patchifier.num_patches

        if self.do_mask_input:
            self.masking = PatchTSTMasking(config)
        else:
            self.masking = nn.Identity()
        self.encoder = PatchTSTEncoder(config, num_patches=num_patches)

        
        self.post_init()

    def forward(
        self,
        past_values: torch.Tensor,
        past_observed_mask: Optional[torch.Tensor] = None,
        future_values: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, PatchTSTModelOutput]:
        r

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        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
        )

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

        
        scaled_past_values, loc, scale = self.scaler(past_values, past_observed_mask)

        
        patched_values = self.patchifier(scaled_past_values)
        if self.do_mask_input:
            masked_values, mask = self.masking(patched_values)
        else:
            masked_values, mask = self.masking(patched_values), None

        encoder_output = self.encoder(
            patch_input=masked_values, output_hidden_states=output_hidden_states, output_attentions=output_attentions
        )

        if not return_dict:
            outputs = (encoder_output.last_hidden_state, encoder_output.hidden_states, encoder_output.attentions)
            outputs = outputs + (mask, loc, scale, patched_values)
            return tuple(v for v in outputs if v is not None)

        return PatchTSTModelOutput(
            last_hidden_state=encoder_output.last_hidden_state,
            hidden_states=encoder_output.hidden_states,
            attentions=encoder_output.attentions,
            mask=mask,
            loc=loc,
            scale=scale,
            patch_input=patched_values,
        )


class PatchTSTMaskPretrainHead(nn.Module):
    

    def __init__(self, config: PatchTSTConfig):
        super().__init__()
        self.dropout = nn.Dropout(config.dropout)
        self.linear = nn.Linear(config.d_model, config.patch_length)
        self.use_cls_token = config.use_cls_token

    def forward(self, embedding: torch.Tensor) -> torch.Tensor:
        
        embedding = self.linear(self.dropout(embedding))  
        if self.use_cls_token:
            embedding = embedding[:, :, 1:, :]  
        return embedding


@add_start_docstrings(
    "The PatchTST for pretrain model.",
    PATCHTST_START_DOCSTRING,
)
class PatchTSTForPretraining(PatchTSTPreTrainedModel):
    def __init__(self, config: PatchTSTConfig):
        super().__init__(config)

        config.do_mask_input = True
        self.model = PatchTSTModel(config=config)
        self.head = PatchTSTMaskPretrainHead(config)

        
        self.post_init()

    def forward(
        self,
        past_values: torch.Tensor,
        past_observed_mask: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, PatchTSTForPretrainingOutput]:
        r

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        
        
        model_output = self.model(
            past_values=past_values,
            past_observed_mask=past_observed_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=True,
        )

        
        
        x_hat = self.head(model_output.last_hidden_state)

        
        loss = nn.MSELoss(reduction="none")
        loss_val = loss(x_hat, model_output.patch_input)
        masked_loss = (loss_val.mean(dim=-1) * model_output.mask).sum() / (model_output.mask.sum() + 1e-10)

        encoder_states = model_output.hidden_states
        if not return_dict:
            outputs = (x_hat,) + model_output[1:-4]
            outputs = (masked_loss,) + outputs if masked_loss is not None else outputs
            return outputs
        return PatchTSTForPretrainingOutput(
            loss=masked_loss, prediction_output=x_hat, hidden_states=encoder_states, attentions=model_output.attentions
        )


class PatchTSTClassificationHead(nn.Module):
    def __init__(self, config: PatchTSTConfig):
        super().__init__()
        self.use_cls_token = config.use_cls_token
        self.pooling_type = config.pooling_type
        self.flatten = nn.Flatten(start_dim=1)
        self.dropout = nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity()
        self.linear = nn.Linear(config.num_input_channels * config.d_model, config.num_targets)

    def forward(self, embedding: torch.Tensor):
        
        if self.use_cls_token:
            
            pooled_embedding = embedding[:, :, 0, :]
        elif self.pooling_type == "mean":
            
            pooled_embedding = embedding.mean(dim=2)
        elif self.pooling_type == "max":
            
            pooled_embedding = embedding.max(dim=2)
        else:
            raise ValueError(f"pooling operator {self.pooling_type} is not implemented yet")
        
        pooled_embedding = self.flatten(pooled_embedding)
        
        output = self.linear(self.dropout(pooled_embedding))
        return output


@add_start_docstrings(
    "The PatchTST for classification model.",
    PATCHTST_START_DOCSTRING,
)
class PatchTSTForClassification(PatchTSTPreTrainedModel):
    def __init__(self, config: PatchTSTConfig):
        super().__init__(config)

        
        if config.do_mask_input:
            logger.warning("Setting `do_mask_input` parameter to False.")
            config.do_mask_input = False

        self.model = PatchTSTModel(config)
        self.head = PatchTSTClassificationHead(config)

        
        self.post_init()

    def forward(
        self,
        past_values: torch.Tensor,
        target_values: torch.Tensor = None,
        past_observed_mask: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[tuple, PatchTSTForClassificationOutput]:
        r

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        model_output = self.model(
            past_values=past_values,
            past_observed_mask=past_observed_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=True,
        )
        y_hat = self.head(model_output.last_hidden_state)

        loss_val = None
        if target_values is not None:
            loss = nn.CrossEntropyLoss()
            loss_val = loss(y_hat, target_values)

        if not return_dict:
            outputs = (y_hat,) + model_output[1:-3]
            outputs = (loss_val,) + outputs if loss_val is not None else outputs
            return outputs
        return PatchTSTForClassificationOutput(
            loss=loss_val,
            prediction_logits=y_hat,
            hidden_states=model_output.hidden_states,
            attentions=model_output.attentions,
        )


@add_start_docstrings(
    "The PatchTST for regression Model.",
    PATCHTST_START_DOCSTRING,
)
class PatchTSTPredictionHead(nn.Module):
    def __init__(self, config: PatchTSTConfig, num_patches, distribution_output=None):
        super().__init__()

        self.share_projection = config.share_projection
        self.num_input_channels = config.num_input_channels
        self.use_cls_token = config.use_cls_token
        self.pooling_type = config.pooling_type
        if self.pooling_type or self.use_cls_token:
            head_dim = config.d_model
        else:
            head_dim = config.d_model * num_patches

        if not self.share_projection:
            
            self.projections = nn.ModuleList()
            self.dropouts = nn.ModuleList()
            self.flattens = nn.ModuleList()
            for i in range(self.num_input_channels):
                self.flattens.append(nn.Flatten(start_dim=2))
                if distribution_output is None:
                    
                    self.projections.append(nn.Linear(head_dim, config.prediction_length))
                else:
                    
                    self.projections.append(distribution_output.get_parameter_projection(head_dim))
                self.dropouts.append(nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity())
        else:
            
            self.flatten = nn.Flatten(start_dim=2)
            if distribution_output is None:
                
                self.projection = nn.Linear(head_dim, config.prediction_length)
            else:
                
                self.projection = distribution_output.get_parameter_projection(head_dim)
            self.dropout = nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity()

    def forward(self, embedding: torch.Tensor):
        
        if self.use_cls_token:
            
            pooled_embedding = embedding[:, :, 0, :]
        else:
            if self.pooling_type == "mean":
                
                pooled_embedding = embedding.mean(dim=2)
            elif self.pooling_type == "max":
                
                pooled_embedding = embedding.max(dim=2)
            else:
                
                pooled_embedding = embedding

        if not self.share_projection:
            output = []
            for i in range(self.num_input_channels):
                
                pooled_embedding = self.flattens[i](pooled_embedding[:, i, :])
                pooled_embedding = self.dropouts[i](pooled_embedding)
                
                
                pooled_embedding = self.projections[i](pooled_embedding)
                output.append(pooled_embedding)
            
            output = torch.stack(output, dim=1)
        else:
            
            pooled_embedding = self.flatten(pooled_embedding)
            pooled_embedding = self.dropout(pooled_embedding)
            
            
            output = self.projection(pooled_embedding)

        if isinstance(output, tuple):
            
            output = tuple(z.transpose(2, 1) for z in output)
        else:
            output = output.transpose(2, 1)  
        return output


@add_start_docstrings(
    "The PatchTST for prediction model.",
    PATCHTST_START_DOCSTRING,
)
class PatchTSTForPrediction(PatchTSTPreTrainedModel):
    def __init__(self, config: PatchTSTConfig):
        super().__init__(config)

        
        if config.do_mask_input:
            logger.warning("Setting `do_mask_input` parameter to False.")
            config.do_mask_input = False

        self.model = PatchTSTModel(config)

        if config.loss == "mse":
            self.distribution_output = None
        else:
            if config.distribution_output == "student_t":
                self.distribution_output = StudentTOutput(dim=config.prediction_length)
            elif config.distribution_output == "normal":
                self.distribution_output = NormalOutput(dim=config.prediction_length)
            elif config.distribution_output == "negative_binomial":
                self.distribution_output = NegativeBinomialOutput(dim=config.prediction_length)
            else:
                raise ValueError(f"Unknown distribution output {config.distribution_output}")

        self.head = PatchTSTPredictionHead(
            config, self.model.patchifier.num_patches, distribution_output=self.distribution_output
        )

        
        self.post_init()

    def forward(
        self,
        past_values: torch.Tensor,
        past_observed_mask: Optional[torch.Tensor] = None,
        future_values: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, PatchTSTForPredictionOutput]:
        r

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        
        model_output = self.model(
            past_values=past_values,
            past_observed_mask=past_observed_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=True,
        )
        
        y_hat = self.head(model_output.last_hidden_state)

        loss_val = None

        if self.distribution_output:
            y_hat_out = y_hat
        else:
            y_hat_out = y_hat * model_output.scale + model_output.loc

        if future_values is not None:
            if self.distribution_output:
                distribution = self.distribution_output.distribution(
                    y_hat, loc=model_output.loc, scale=model_output.scale
                )
                loss_val = nll(distribution, future_values)
                
                loss_val = weighted_average(loss_val)
            else:
                loss = nn.MSELoss(reduction="mean")
                loss_val = loss(y_hat_out, future_values)

        loc = model_output.loc
        scale = model_output.scale

        if not return_dict:
            outputs = (y_hat_out,) + model_output[1:-1]
            outputs = (loss_val,) + outputs if loss_val is not None else outputs
            return outputs
        return PatchTSTForPredictionOutput(
            loss=loss_val,
            prediction_outputs=y_hat_out,
            hidden_states=model_output.hidden_states,
            attentions=model_output.attentions,
            loc=loc,
            scale=scale,
        )

    def generate(
        self,
        past_values: torch.Tensor,
        past_observed_mask: Optional[torch.Tensor] = None,
    ) -> SamplePatchTSTOutput:
        
        
        num_parallel_samples = self.config.num_parallel_samples

        
        outputs = self(
            past_values=past_values,
            future_values=None,
            past_observed_mask=past_observed_mask,
            output_hidden_states=False,
        )
        if self.distribution_output:
            
            distribution = self.distribution_output.distribution(
                outputs.prediction_outputs, loc=outputs.loc, scale=outputs.scale
            )
            
            samples = [distribution.sample() for _ in range(num_parallel_samples)]
            
            samples = torch.stack(samples, dim=1)
        else:
            samples = outputs.prediction_outputs.unsqueeze(1)

        return SamplePatchTSTOutput(sequences=samples)


class PatchTSTRegressionHead(nn.Module):
    

    def __init__(self, config: PatchTSTConfig, distribution_output=None):
        super().__init__()
        self.y_range = config.output_range
        self.use_cls_token = config.use_cls_token
        self.pooling_type = config.pooling_type
        self.distribution_output = distribution_output

        head_dim = config.num_input_channels * config.d_model

        self.flatten = nn.Flatten(start_dim=1)
        self.dropout = nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity()

        if distribution_output is None:
            self.projection = nn.Linear(head_dim, config.num_targets)
        else:
            self.projection = distribution_output.get_parameter_projection(head_dim)

    def forward(self, embedding: torch.Tensor):
        
        if self.use_cls_token:
            
            pooled_embedding = embedding[:, :, 0, :]
        elif self.pooling_type == "mean":
            
            pooled_embedding = embedding.mean(dim=2)
        elif self.pooling_type == "max":
            
            pooled_embedding = embedding.max(dim=2)
        else:
            raise ValueError(f"pooling operator {self.pooling_type} is not implemented yet")
        
        
        pooled_embedding = self.dropout(self.flatten(pooled_embedding))
        
        
        output = self.projection(pooled_embedding)
        
        if (self.distribution_output is None) & (self.y_range is not None):  
            output = torch.sigmoid(output) * (self.y_range[1] - self.y_range[0]) + self.y_range[0]
        return output


@add_start_docstrings(
    "The PatchTST for regression model.",
    PATCHTST_START_DOCSTRING,
)
class PatchTSTForRegression(PatchTSTPreTrainedModel):
    def __init__(self, config: PatchTSTConfig):
        super().__init__(config)

        
        if config.do_mask_input:
            logger.warning("Setting `do_mask_input` parameter to False.")
            config.do_mask_input = False

        self.model = PatchTSTModel(config)
        if config.loss == "mse":
            self.distribution_output = None
        else:
            if config.distribution_output == "student_t":
                self.distribution_output = StudentTOutput(dim=config.prediction_length * config.num_targets)
            elif config.distribution_output == "normal":
                self.distribution_output = NormalOutput(dim=config.prediction_length * config.num_targets)
            elif config.distribution_output == "negative_binomial":
                self.distribution_output = NegativeBinomialOutput(dim=config.prediction_length * config.num_targets)
            else:
                raise ValueError(f"Unknown distribution output {config.distribution_output}")

        self.head = PatchTSTRegressionHead(config, self.distribution_output)

        
        self.post_init()

    def forward(
        self,
        past_values: torch.Tensor,
        target_values: torch.Tensor = None,
        past_observed_mask: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[tuple, PatchTSTForRegressionOutput]:
        r

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        model_output = self.model(
            past_values=past_values,
            past_observed_mask=past_observed_mask,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=True,
        )
        
        y_hat = self.head(model_output.last_hidden_state)

        loss = None
        if target_values is not None:
            if self.distribution_output:
                distribution = self.distribution_output.distribution(y_hat)
                loss = nll(distribution, target_values)
                
                loss = weighted_average(loss)
            else:
                loss = nn.MSELoss(reduction="mean")
                loss = loss(y_hat, target_values)

        if not return_dict:
            outputs = (y_hat,) + model_output[1:-3]
            outputs = (loss,) + outputs if loss is not None else outputs
            return outputs
        return PatchTSTForRegressionOutput(
            loss=loss,
            regression_outputs=y_hat,
            hidden_states=model_output.hidden_states,
            attentions=model_output.attentions,
        )

    def generate(
        self,
        past_values: torch.Tensor,
        past_observed_mask: Optional[torch.Tensor] = None,
    ) -> SamplePatchTSTOutput:
        
        
        num_parallel_samples = self.config.num_parallel_samples

        
        outputs = self(
            past_values=past_values,
            target_values=None,
            past_observed_mask=past_observed_mask,
            output_hidden_states=False,
        )

        
        distribution = self.distribution_output.distribution(outputs.regression_outputs)
        
        samples = [distribution.sample() for _ in range(num_parallel_samples)]
        
        samples = torch.stack(samples, dim=1)
        return SamplePatchTSTOutput(sequences=samples)