o Zh7Q @sdZddlZddlmZddlmZmZmZddlZddl m Z ddl m Z ddl mZddlmZmZmZdd lmZmZdd lmZd d lmZeeZGd dde jZGddde jZGddde jZ Gddde jZ!Gddde jZ"Gddde jZ#Gddde jZ$Gddde jZ%Gddde jZ&Gdd d e jZ'Gd!d"d"e jZ(Gd#d$d$e jZ)Gd%d&d&e jZ*eGd'd(d(e Z+Gd)d*d*e jZ,  + djd,ej-d-e.d.ee/d/e0d0e1f d1d2Z2  dkd,ej-d3ee/e1fd.ee/d0e1fd4d5Z3Gd6d7d7e jZ4Gd8d9d9e jZ5Gd:d;d;e jZ6Gdd?d?e jZ8eGd@dAdAeZ9GdBdCdCe+Z:eGdDdEdEeZ;edFdGGdHdIdIe+ZeGdNdOdOeZ?eGdPdQdQeZ@eGdRdSdSeZAdTejBjCdUej-dVej-fdWdXZDdldYej-dZeej-dVej-fd[d\ZEGd]d^d^e+ZFeGd_d`d`eZGGdadbdbe+ZHeGdcddddeZIGdedfdfe jZJGdgdhdhe+ZKgdiZLdS)mzPyTorch PatchTSMixer model.N) dataclass)OptionalTupleUnion)PreTrainedModel) ModelOutput)NegativeBinomialOutput NormalOutputStudentTOutput)auto_docstringlogging)deprecate_kwarg)PatchTSMixerConfigcs2eZdZdZdedeffdd ZddZZS)PatchTSMixerGatedAttentionz Module that applies gated attention to input data. Args: in_size (`int`): The input size. out_size (`int`): The output size. in_sizeout_sizecs*tt|||_tjdd|_dS)Ndim)super__init__nnLinear attn_layerZSoftmax attn_softmax)selfrr __class__e/var/www/auris/lib/python3.10/site-packages/transformers/models/patchtsmixer/modeling_patchtsmixer.pyr-s z#PatchTSMixerGatedAttention.__init__cCs|||}||}|SN)rr)rinputsZ attn_weightr r r!forward2sz"PatchTSMixerGatedAttention.forward)__name__ __module__ __qualname____doc__intrr$ __classcell__r r rr!r$src6eZdZdZdeffdd ZdejfddZZ S)PatchTSMixerBatchNormzP Compute batch normalization over the sequence length (time) dimension. configcs"ttj|j|jd|_dS)NZeps)rrr BatchNorm1dd_modelnorm_eps batchnormrr-rr r!r>s zPatchTSMixerBatchNorm.__init__r#cCs"|dd}||}|ddS)a 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)` r) transposer2)rr#outputr r r!r$Bs   zPatchTSMixerBatchNorm.forward r%r&r'r(rrtorchTensorr$r*r r rr!r,9sr,csNeZdZdZdeffdd ZededejfddZ de j fd d Z Z S) PatchTSMixerPositionalEncodingz' Class for positional encoding r-cs<t|jr|||_dStt|j |j |_dSr") rruse_positional_encoding_init_pe position_encr Parameterr8zeros num_patchesr0r3rr r!rTs z'PatchTSMixerPositionalEncoding.__init__returncCs|jdkrtjt|j|jdd}|S|jdkrvt|j|j}td|j d}t td|jdt d|j }t |||dddddf<t|||dddddf<||}||d }tj|d d}|St|jd ) NrandomTZ requires_gradZsincosrrr4g@ FzN is not a valid positional encoder. Available types are 'random' and 'sincos'.)positional_encoding_typerr>r8Zrandnr@r0r?Zarange unsqueezeexpmathlogsincosmeanstd ValueError)r-r=positionZdiv_termr r r!r<\s  (    z'PatchTSMixerPositionalEncoding._init_pe patch_inputcCs||j}|Sr")r=)rrP hidden_stater r r!r$ps z&PatchTSMixerPositionalEncoding.forward)r%r&r'r(rr staticmethodrr>r<r8r9r$r*r r rr!r:Os r:cr+)PatchTSMixerNormLayerzeNormalization block Args: config (`PatchTSMixerConfig`): Configuration. r-csFt|j|_d|jvrt||_dStj|j|j d|_dS)Nbatchr.) rrnorm_mlplowerr,normr LayerNormr0r1r3rr r!r~s zPatchTSMixerNormLayer.__init__r#cCsfd|jvr,t||jd|jd|jd|jdf}||}t||j}|S||}|S)a Args: inputs (`torch.Tensor` of shape `((batch_size, num_channels, num_patches, d_model))`): Input to the normalization layer. Returns: `torch.Tensor` of shape `((batch_size, num_channels, num_patches, d_model))` rTrrr4r)rUrVr8reshapeshaperW)rr#Zinputs_reshapedr r r!r$s  zPatchTSMixerNormLayer.forwardr7r r rr!rSvs rScs,eZdZfddZdejfddZZS)PatchTSMixerMLPcsPt||j}t|||_t|j|_t|||_ t|j|_ dSr") rrZexpansion_factorrrfc1Dropoutdropoutdropout1fc2dropout2)r in_features out_featuresr-Z num_hiddenrr r!rs  zPatchTSMixerMLP.__init__r#cCs0|tj||}||}||}|S)z Args: inputs (`torch.Tensor` of shape `((batch_size, num_channels, num_patches, d_model))`): Input to the MLP layer. Returns: `torch.Tensor` of the same shape as `inputs` )r_r functionalZgelur\r`ra)rr#r r r!r$s  zPatchTSMixerMLP.forward)r%r&r'rr8r9r$r*r r rr!r[s r[cr+)$PatchTSMixerChannelFeatureMixerBlockzThis module mixes the features in the channel dimension. Args: config (`PatchTSMixerConfig`): Configuration. r-cPtt||_|j|_t|j|j|d|_|jr&t|j|jd|_ dSdSNrbrcr-rr) rrrSrW gated_attnr[num_input_channelsmlpr gating_blockr3rr r!rs   z-PatchTSMixerChannelFeatureMixerBlock.__init__r#cCsT|}||}|dddd}|jr||}||}|dddd}||}|S)z Args: inputs (`torch.Tensor` of shape `((batch_size, num_channels, num_patches, d_model))`): input to the MLP layer Returns: `torch.Tensor` of the same shape as `inputs` rrr4r)rWZpermuterjrmrl)rr#residualoutr r r!r$s   z,PatchTSMixerChannelFeatureMixerBlock.forwardr7r r rr!resrecseZdZdZ      ddededed ed ed ed eed eeffdd Z e ddde ddde ddd      dde j dee j dee e j dee j dee j dedee j de e j ee j ee e j ffddZZS)PatchTSMixerAttentionz=Multi-headed attention from 'Attention Is All You Need' paperFTN embed_dim num_headsr^ is_decoderbias is_causalr- layer_idxc st||_||_||_|||_||_|j||jkr*td|jd|d|jd|_||_ ||_ ||_ |durK|j rKt d|jjdtj|||d|_tj|||d|_tj|||d|_tj|||d|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: ).gzInstantiating a decoder z without passing `layer_idx` is not recommended and will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` when creating this class.)ru)rrrrrsr^head_dimr-rNscalingrtrvrwloggerZ warning_oncerr%rrk_projv_projq_projout_proj) rrrrsr^rtrurvr-rwrr r!rs0    zPatchTSMixerAttention.__init__key_value_statesz4.55)versionpast_key_valuecache_position hidden_statesattention_masklayer_head_maskoutput_attentionsrAcCs|\}} } |||d|j|jdd} | |j} |||d|j|jdd} |||d|j|jdd} ||jd|jf}| j |} | j |} | j |} | d}t | | dd}|||j| |fkrt d||j| |fd||dur|ddddddd| j df}|||j| ||}|||j| |}tjj|dd}|dur||jfkrt d |jfd||dddd|||j| |}|||j| |}|r|||j| |}|||j| |}nd}tjj||j|jd }t || }|||j| |jfkr8t d ||j| |jfd||||j| |j}|dd}| || |j}||}||dfS) z#Input shape: Batch x Time x Channelrrr4z$Attention weights should be of size z , but is Nrz/Head mask for a single layer should be of size )ptrainingz `attn_output` should be of size )sizer~viewrsryr5rzr|r}rYr8ZbmmrNrZrrdZsoftmaxr^rrrr)rrrrrrrrZbszZtgt_len_Z query_statesZ key_statesZ value_statesZ proj_shapeZsrc_lenZ attn_weightsZattn_weights_reshapedZ attn_probsZ attn_outputr r r!r$s`" ""    & "    zPatchTSMixerAttention.forward)rqFTFNN)NNNNFN)r%r&r'r(r)floatboolrrrrr8r9rr$r*r r rr!rpsf  (    rpc.eZdZdZdeffdd ZddZZS)PatchMixerBlockzxThis module mixes the patch dimension. Args: config (`PatchTSMixerConfig`): Configuration. r-cs~tt||_|j|_|j|_t|j|j|d|_|jr(t |j|jd|_ |jr=t |j |j |jd|_t||_dSdS)Nrhri)rrrsr^)rrrSrW self_attnrjr[r@rlrrmrpr0Zself_attn_headsr^self_attn_layer norm_attnr3rr r!rxs&  zPatchMixerBlock.__init__c Cs|}||}|jr,|j\}}}}|||||}|j|dd\}} } |||||}|dd}||}|jr?||}|dd}|jrO| ||}||} | S)z Args: hidden_state (`torch.Tensor`): Input tensor. Returns: `torch.Tensor`: Transformed tensor. F)rr4r) rWrrZrYrr5rlrjrmr) rrQrn batch_sizeZn_varsr@r0Zhidden_state_reshapedZx_attnrror r r!r$s      zPatchMixerBlock.forwardr%r&r'r(rrr$r*r r rr!rpsrcr+)FeatureMixerBlockzThis module mixes the hidden feature dimension. Args: config (`PatchTSMixerConfig`): Configuration. r-crfrg) rrrSrWrjr[r0rlrrmr3rr r!rs  zFeatureMixerBlock.__init__hiddencCs4|}||}||}|jr||}||}|S) Args: hidden (`torch.Tensor` of shape `(batch_size, num_patches, d_model)`): Input tensor to the layer. Returns: `torch.Tensor`: Transformed tensor. )rWrlrjrm)rrrnror r r!r$s   zFeatureMixerBlock.forwardr7r r rr!rsrcr+)PatchTSMixerLayerz The `PatchTSMixer` layer that does all three kinds of mixing. Args: config (`PatchTSMixerConfig`): Configuration. r-csHtt|d|_t|d|_|j|_|jdkr"t|d|_dSdS)Nr- mix_channel) rrr patch_mixerr feature_mixermoderechannel_feature_mixerr3rr r!rs    zPatchTSMixerLayer.__init__rcCs,|jdkr ||}||}||}|S)rr)rrrr)rrr r r!r$s   zPatchTSMixerLayer.forwardr7r r rr!rs  rcs6eZdZdZdeffdd Zd defddZZS) PatchTSMixerBlockzThe main computing framework of the `PatchTSMixer` model. Args: config (`PatchTSMixerConfig`): Configuration. r-cs2tj}tfddt|D|_dS)Ncsg|]}tdqS)r)r.0rrr r! sz.PatchTSMixerBlock.__init__..)rr num_layersrZ ModuleListrangemixers)rr-rrrr!rs "zPatchTSMixerBlock.__init__Foutput_hidden_statescCs>g}|}|jD] }||}|r||q|r||fS|dfS)as Args: hidden_state (`torch.Tensor`): The input tensor. output_hidden_states (`bool`, *optional*, defaults to False.): Whether to output the hidden states as well. Returns: `torch.Tensor`: The embedding. `list`: List of all hidden states if `output_hidden_states` is set to `True`. N)rappend)rrQrZall_hidden_statesZ embeddingmodr r r!r$s   zPatchTSMixerBlock.forwardF) r%r&r'r(rrrr$r*r r rr!rsrc0eZdZdZddeffdd ZddZZS) PatchTSMixerForPredictionHeadzqPrediction Head for Forecasting Args: config (`PatchTSMixerConfig`): Configuration. Nr-cs|t|j|_|jdur|jt|j|_|dur+t|j |j |j |_ n | |j |j |_ tjdd|_dS)NrZ start_dim)rrprediction_channel_indicessortrr] head_dropout dropout_layerrr@r0prediction_lengthbase_forecast_blockget_parameter_projectionFlattenflatten)rr-distribution_outputrr r!r9s    z&PatchTSMixerForPredictionHead.__init__cs|}|}|}t|trtdd|D}n|dd}jdurBt|tr;tfdd|D}|S|djf}|S)ar Args: hidden_features (`torch.Tensor` of shape `(batch_size, num_patch, d_model)` in `flatten` mode or `(batch_size, n_vars, num_patch, d_model)` in `common_channel`/`mix_channel` mode.): Input hidden features. Returns: `torch.Tensor` of shape `(batch_size, prediction_length, nvars)`. css|] }|ddVqdS)rrN)r5rzr r r! \sz8PatchTSMixerForPredictionHead.forward..rrNc3s|] }|djfVqdS).N)rrrr r!rbs.)rrr isinstancetupler5rrhidden_featuresforecastr rr!r$Ks      z%PatchTSMixerForPredictionHead.forwardr"rr r rr!r1srcr) PatchTSMixerLinearHeadzLinear head for Classification and Regression. Args: config (`PatchTSMixerConfig`): Configuration. Nr-cst|j|_|j|_|jdur|j}nd}||_|dur.t|j|j ||j |_ n | |j|j ||_ |jdurGtj dd|_ntj dd|_t|j|_dS)Nrrr)rrhead_aggregation output_ranger@rrrr0rk num_targets projectionrrrr]rr^)rr-rZ mul_factorrr r!rqs&   zPatchTSMixerLinearHead.__init__cCs|dd}|jdkr|d}n|jdkr|jddj}n |jdkr(|jdd}|jr0||}||}||}|jdurX|j durXt ||j d |j d |j d }|S) ai Args: hidden_features (`torch.Tensor` of shape `(batch_size x num_patch x d_model)` in `flatten` mode or `(batch_size x n_vars x num_patch x d_model)` in `common_channel`/`mix_channel` mode.): Input hidden features. Returns: `torch.Tensor` of shape `(batch_size x num_targets)`. rrZuse_last).rZmax_poolrZavg_poolNrr) r5rmaxvaluesrLrr^rrrr8Zsigmoid)rrr r r!r$s        &zPatchTSMixerLinearHead.forwardr"rr r rr!risrc@s$eZdZeZdZdZdZddZdS)PatchTSMixerPreTrainedModelmodel past_valuesFcCst|tr|jjdkrtjj|jddddSdSt|tjtj fr1|j j |j j ddSt|trG|jj j |jj j ddSt|tjre|j j jd|jjd|j durg|j j dSdSdS)zInitialize weightsrBrqg?)rLrM?N)rr:r-rErinitZnormal_r=rXr/rudataZzero_weightZfill_r,r2rZinit_std)rmoduler r r! _init_weightss      z)PatchTSMixerPreTrainedModel._init_weightsN) r%r&r'rZ config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrr r r r!rs  rcr)PatchTSMixerPretrainHeadzcPretraining head. Args: config (`PatchTSMixerConfig`): Configuration. r-cs.tt|j|_t|j|j|_ dSr") rrrr]rrrr0 patch_length base_pt_blockr3rr r!rs z!PatchTSMixerPretrainHead.__init__cCs||}||}|S)a Args: hidden_features (`torch.Tensor` of shape `(batch_size x num_patch x d_model)` in `flatten` mode or `(batch_size x n_vars x num_patch x d_model)` in `common_channel`/`mix_channel` mode.): Input hidden features. Returns: `torch.Tensor` of shape `(batch_size x n_vars x num_patch x patch_length)`. )rrrr r r!r$s z PatchTSMixerPretrainHead.forwardrr r rr!rsrFr# mask_ratiounmasked_channel_indiceschannel_consistent_masking mask_valuecCs*|dks|dkrtd|d|j\}}}}|j} t|d|} |r5tj|d|| d} | d|d} n tj|||| d} tj|||| d} d| ddddd| f<tj| dd} tj| dd}tj | d|d } | dddd|} |durd| dd|ddddf<| | |}|| d fS) arandom_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] rrz Mask ratio z has to be between 0 and 1.deviceNrr)rindex.r) rNrZrr)r8ZrandrepeatZonesZargsortgatherrF masked_fillr)r#rrrrr num_channelssequence_length num_featuresrZlen_keepnoisemaskZ ids_shuffleZ ids_restore inputs_maskr r r!random_maskings& rnum_forecast_mask_patchescCst|tr|g}dd|D}|j\}}}}tj||||jd} g} d} t|} t||D](\} }| dks9| |krAtd| dt||| }| | ||g| |7} q-t | ddd } | |krq| dd || | dd <n| |kr| d d | || d d <d}| D]\}}}||}d | ||d d | d f<|}qt | jd}| |} | d  d d d |} |d urd| d d |d d d d f<|| |}|| dfS)aForecast 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_length)` 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)` cSsg|]}dqS)rr rr r r!r>sz$forecast_masking..rrznum_forecast_mask_patches z6 should be greater than 0 and less than total patches.cSs|dSNr4r )xr r r!Psz"forecast_masking..)keyr4rrNr)rr)rZr8r?rsumziprNrsortedZrandpermrFrrr)r#rrrZforecast_mask_ratiosrrrrrZt_list total_length total_ratiorratioZtemp_lenZbatch1Z patch_lenrZbatch2permrr r r!forecast_masking$sB    rcr+)PatchTSMixerPatchifyz A class to patchify the time series sequence into different patches Returns: `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)` r-cst|j|_|j|_|j|_|j|jkr$td|jd|jdt|j|j|j|jd|_|j|j|jd}|j||_ dS)NzSequence length (z+) has to be greater than the patch length ()r) rrZcontext_lengthrr patch_striderNrr@sequence_start)rr-Znew_sequence_lengthrr r!rqs   zPatchTSMixerPatchify.__init__rcCsp|jd}||jkrtd|d|jd|dd|jdddf}|jd|j|jd}|dd}|S)a! 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)` rzInput sequence length (z%) doesn't match model configuration (rxN) dimensionrstepr) rZrrNrZunfoldrrr5 contiguous)rrrr6r r r!r$s zPatchTSMixerPatchify.forwardr7r r rr!risrcr+)PatchTSMixerMaskinga Class to perform random or forecast masking. Parameters: config (`PatchTSMixerConfig`): 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 r-csXt|j|_|j|_|j|_|j|_|j|_|j|_|jdur*t|j|_dSdSr") rrrandom_mask_ratior mask_typerrrrr3rr r!rs  zPatchTSMixerMasking.__init__rPcCsr|jdkrt||j|j|j|jd\}}n|jdkr(t||j|j|jd\}}n td|jd| }||fS)a 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 rB)r#rrrrr)r#rrrzInvalid mask type .) rrrrrrrrrNr)rrPZ masked_inputrr r r!r$s$   zPatchTSMixerMasking.forwardr7r r rr!rs  rc PeZdZdZdeffdd Zdejdejdeejejejffdd Z Z S) PatchTSMixerStdScalerz 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. r-csVtt|dr |jnd|_t|dr|jnd|_t|dr&|j|_dSd|_dS)N scaling_dimrkeepdimT minimum_scalegh㈵>)rrhasattrrrrrr3rr r!rs  zPatchTSMixerStdScaler.__init__robserved_indicatorrAcCsz|j|j|jd}|d}||j|j|jd|}|||dj|j|jd|}t||j}|||||fS)C 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)`) rrr4)rrrZ clamp_minr8sqrtr)rrr denominatorlocZvariancescaler r r!r$s  "zPatchTSMixerStdScaler.forward r%r&r'r(rrr8r9rr$r*r r rr!rsrc r) PatchTSMixerMeanScalerz Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data accordingly. r-csltt|dr |jnd|_t|dr|jnd|_t|dr#|jnd|_t|dr1|j|_dSd|_dS)NrrrTr绽|= default_scale)rrrrrrrrr3rr r!rs  zPatchTSMixerMeanScaler.__init__rrrAc Cs||j|jdd}|j|jdd}|tj|dd}|jdur:|jdd}tj|ddd}t||}n|jt|}t|dk||}tj||j d}||} |j sa|j|jd}| t ||fS)rTrrminNrr) absrrr8clamprZsqueeze ones_likewhererr zeros_like) rrrZts_sumZ num_observedr Z batch_sumZbatch_observationsrZ scaled_datar r r!r$s  zPatchTSMixerMeanScaler.forwardr r r rr!rsrc sXeZdZdZdeffdd Z d dejdeejde ejejejffd d Z Z S) PatchTSMixerNOPScalerz| Assigns a scaling factor equal to 1 along the first dimension, and therefore applies no scaling to the input data. r-cs@tt|dr |jnd|_t|dr|j|_dSd|_dS)NrrrT)rrrrrrr3rr r!r7s  zPatchTSMixerNOPScaler.__init__NrrrAcCsBtj|ddj|j|jd}tj|ddj|j|jd}|||fS)a 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)`) FrC)rr)r8rrLrrr)rrrr r r r r!r$<s  zPatchTSMixerNOPScaler.forwardr") r%r&r'r(rrr8r9rrr$r*r r rr!r2src@s:eZdZUdZdZeejed<dZ ee ejed<dS)PatchTSMixerEncoderOutputa Base class for `PatchTSMixerEncoderOutput`, with potential hidden states. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, d_model)`): Hidden-state at the output of the last layer of the model. hidden_states (`tuple(torch.FloatTensor)`, *optional*): Hidden-states of the model at the output of each layer. Nlast_hidden_stater) r%r&r'r(rrr8 FloatTensor__annotations__rrr r r r!rMs  rc s\eZdZdZdeffdd Ze  d dejde e d e e d e e e ffd d ZZS)PatchTSMixerEncoderz Encoder for PatchTSMixer which inputs patched time-series and outputs patched embeddings. Args: config (`PatchTSMixerConfig`): Configuration. r-csbt||j|_t|j|j|_|jrt |d|_ nd|_ t |d|_ |j r/| dSdS)Nr)rruse_return_dictrrrr0patcherr;r:positional_encoderrmlp_mixer_encoder post_initr3rr r!rfs   zPatchTSMixerEncoder.__init__FNrr return_dictrAcCsh|dur|n|j}||}|jdur||}|j||d\}}|s.tdd||fDSt||dS)a past_values (`torch.FloatTensor` of shape `(batch_size, seq_length, num_input_channels)`): Context values of the time series. For a pretraining task, this denotes the input time series to predict the masked portion. For a forecasting task, this denotes the history/past time series values. Similarly, for classification or regression tasks, it denotes the appropriate context values of the time series. For univariate time series, `num_input_channels` dimension should be 1. For multivariate time series, it is greater than 1. Returns: `torch.FloatTensor` of shape `(batch_size, n_vars, num_patches, d_model)` N)rcs|]}|VqdSr"r rvr r r!r  z.PatchTSMixerEncoder.forward..)rr)rrr r!rr)rrrr#Zpatchesrrr r r!r$vs    zPatchTSMixerEncoder.forward)FN)r%r&r'r(rrr r8r9rrrrrr$r*r r rr!r]s rc@seZdZUdZdZeejed<dZ ee ejed<dZ eejed<dZ eejed<dZ eejed<dZeejed<dS) PatchTSMixerModelOutputa Base class for model's outputs, with potential hidden states. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, d_model)`): Hidden-state at the output of the last layer of the model. hidden_states (`tuple(torch.FloatTensor)`, *optional*): Hidden-states of the model at the output of each layer. patch_input (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`): Patched input data to the model. mask: (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches)`,*optional*): Bool Tensor indicating True in masked patches and False otherwise. loc: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`,*optional*): Gives the mean of the context window per channel. Used for revin denorm outside the model, if revin enabled. scale: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`,*optional*): Gives the std dev of the context window per channel. Used for revin denorm outside the model, if revin enabled. NrrrPrr r )r%r&r'r(rrr8rrrrrPrr r r r r r!r(s r(z= The PatchTSMixer Model for time-series forecasting. )Z custom_introc sbeZdZddedeffdd Ze   ddejde ejd e ed e ed e f d d Z Z S)PatchTSMixerModelFr- mask_inputcst||j|_t||_t||_|durt||_nd|_|j dkr,t ||_ n|j dks6|j dur.)rrrPrr r ) rr8rr.r,r-r+rrrrrr() rrr/rr#rZscaled_past_valuesr r Z patched_xZ enc_inputZencoder_outputr r r!r$sD     zPatchTSMixerModel.forwardr)NFN) r%r&r'rrrr r8r9rr(r$r*r r rr!r)s"r)c@^eZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dS) PatchTSMixerForPreTrainingOutputa Output type of [`PatchTSMixerForPreTrainingOutput`]. Args: prediction_outputs (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, patch_length)`): Prediction output from the pretrain head. hidden_states (`tuple(torch.FloatTensor)`, *optional*): Hidden-states of the model at the output of each layer. last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, d_model)`): Backbone embeddings before passing through the head. loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`): Total loss Nlossprediction_outputsrr r%r&r'r(r3rr8rrr4rrrr r r r!r2& r2cfeZdZdZdeffdd Ze    ddejd e ejd e e d e d e e d e f ddZ Z S)PatchTSMixerForPretrainingz `PatchTSMixer` for mask pretraining. Args: config (`PatchTSMixerConfig`): Configuration. Returns: `None`. r-csLt|t|dd|_t|d|_|j|_|j|_|jr$|dSdS)NT)r*r) rrr)rrhead masked_lossrr"r3rr r!rHs   z#PatchTSMixerForPretraining.__init__NFTrr/r return_lossr#rAc Cs|dur|n|j}|jdurtjjdd}ntjjdd}|j||||d}t|tr/t|}| |j }|dur@|||j } nd} |jdur]| dur]| j dd|j |j d } |sntd d | ||j |jfDSt| ||j |jd S) aT past_values (`torch.FloatTensor` of shape `(batch_size, seq_length, num_input_channels)`): Context values of the time series. For a pretraining task, this denotes the input time series to predict the masked portion. For a forecasting task, this denotes the history/past time series values. Similarly, for classification or regression tasks, it denotes the appropriate context values of the time series. For univariate time series, `num_input_channels` dimension should be 1. For multivariate time series, it is greater than 1. observed_mask (`torch.FloatTensor` 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_loss (`bool`, *optional*): Whether to return the loss in the `forward` call. NTnoneZ reductionrLr/rr#rrrcsr$r"r r%r r r!rr'z5PatchTSMixerForPretraining.forward..r3r4rr)rr:r8rMSELossrrrr(r9rrPrLrrrr2) rrr/rr;r#r3 model_outputZx_hatloss_valr r r!r$Ss@   $ z"PatchTSMixerForPretraining.forwardNFTN)r%r&r'r(rrr r8r9rrr2r$r*r r rr!r8<s*  r8c@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dZ eejed<dZeejed<dS) PatchTSMixerForPredictionOutputa Output type of [`PatchTSMixerForPredictionOutput`]. Args: prediction_outputs (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_input_channels)`): Prediction output from the forecast head. last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, d_model)`): Backbone embeddings before passing through the head. hidden_states (`tuple(torch.FloatTensor)`, *optional*): Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`): Total loss. loc (`torch.FloatTensor`, *optional* of shape `(batch_size, 1, num_input_channels)`): Input mean scale (`torch.FloatTensor`, *optional* of shape `(batch_size, 1, num_input_channels)`): Input std dev Nr3r4rrr r )r%r&r'r(r3rr8rrr4rrrr r r r r r!rDs rDc@$eZdZUdZdZeejed<dS)"SamplePatchTSMixerPredictionOutputa9 Base class for time series model's predictions outputs that contains the sampled values from the chosen distribution. Args: sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, prediction_length, number_channels)`): Sampled values from the chosen distribution. N sequences r%r&r'r(rGrr8rrr r r r!rF  rFc@rE)"SamplePatchTSMixerRegressionOutputa$ Base class for time series model's predictions outputs that contains the sampled values from the chosen distribution. Args: sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, num_targets)` Sampled values from the chosen distribution. NrGrHr r r r!rJrIrJinputtargetrAcCs || S)zc Computes the negative log likelihood loss from input distribution with respect to target. )Zlog_prob)rKrLr r r!nlls rM input_tensorweightscCsr|dur3t|dk||t|}tj|r|j|dn|dd}|r-|j|d|S||S|j|dS)aj 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`. Nrrrr)r8rrrrrL)rNrOrZweighted_tensorZ sum_weightsr r r!weighted_averages "  rPcseZdZdZdeffdd Ze     ddejd e ejd e ejd e e d e d e e de fddZ ddejd e ejde fddZZS)PatchTSMixerForPredictionz `PatchTSMixer` for forecasting application. Args: config (`PatchTSMixerConfig`): Configuration. Returns: `None`. r-cst||j|_|j|_|j|_|j|_|jdkrd|_n#|j}tt t d}| |jd}|dur:||d|_nt d|jt ||_t||jd|_|jrX|dSdS)NmseZ student_tnormalnegative_binomialrUnknown distribution output r-r)rrr3rrnum_parallel_samplesrrr r r getrNr)rrr9r")rr-rdistribution_output_map output_classrr r!rs0    z"PatchTSMixerForPrediction.__init__NFTrr/ future_valuesrr;r#rAcCs|jdkr tjdd}n |jdkrt}ntd|dur|n|j}|j||||d}t|tr3t |}| |j } d} |j dur|j ro|j j| |jd|j f|jd|j fd } |durn|d urn|| |d|j f} t| } nZ| |jd|j f|jd|j f} |dur|d ur|| |d|j f} n5|j r|j j| |j|jd } |dur|d ur|| |} t| } n| |j|j} |dur|d ur|| |} |j dur|jd|j f} |jd|j f} n|j} |j} |std d | | |j |j| | fDSt| | |j |j| | d S)a past_values (`torch.FloatTensor` of shape `(batch_size, seq_length, num_input_channels)`): Context values of the time series. For a pretraining task, this denotes the input time series to predict the masked portion. For a forecasting task, this denotes the history/past time series values. Similarly, for classification or regression tasks, it denotes the appropriate context values of the time series. For univariate time series, `num_input_channels` dimension should be 1. For multivariate time series, it is greater than 1. observed_mask (`torch.FloatTensor` 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.FloatTensor` of shape `(batch_size, target_len, num_input_channels)` for forecasting,: `(batch_size, num_targets)` for regression, or `(batch_size,)` for classification, *optional*): Target values of the time series, that serve as labels for the model. The `future_values` is what the Transformer needs during training to learn to output, given the `past_values`. Note that, this is NOT required for a pretraining task. For a forecasting task, the shape is be `(batch_size, target_len, num_input_channels)`. Even if we want to forecast only specific channels by setting the indices in `prediction_channel_indices` parameter, pass the target data with all channels, as channel Filtering for both prediction and target will be manually applied before the loss computation. return_loss (`bool`, *optional*): Whether to return the loss in the `forward` call. rRrLr=rM2Invalid loss function: Allowed values: mse and nllNr>.r r Tcsr$r"r r%r r r!rr'z4PatchTSMixerForPrediction.forward..)r3r4rrr r )r3rr@rMrNrrrrr(r9rrr distributionr r rPrrD)rrr/r\rr;r#r3rAy_hatrBr_r r r r r!r$!s $          z!PatchTSMixerForPrediction.forwardcs\|j}||d|dd}|jj|j|j|jdfddt|D}tj|dd}t |d S) a Generate sequences of sample predictions from a model with a probability distribution head. Args: 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. 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: [`SamplePatchTSMixerPredictionOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of samples, prediction_length, num_input_channels)`. NF)rr\r/rr^cg|]}qSr samplerr_r r!rsz6PatchTSMixerForPrediction.generate..rrrG) rXrr_r4r r rr8stackrF)rrr/rXoutputssamplesr rdr!generates  z"PatchTSMixerForPrediction.generate)NNFTNr")r%r&r'r(rrr r8r9rrrDr$rFrir*r r rr!rQs@   |rQc@r1)-PatchTSMixerForTimeSeriesClassificationOutputa Output type of [`PatchTSMixerForTimeSeriesClassificationOutput`]. Args: prediction_outputs (`torch.FloatTensor` of shape `(batch_size, num_labels)`): Prediction output from the classification head. last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, d_model)`): Backbone embeddings before passing through the head. hidden_states (`tuple(torch.FloatTensor)`, *optional*): Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`): Total loss. Nr3r4rrr5r r r r!rjr6rjcr7)'PatchTSMixerForTimeSeriesClassificationz `PatchTSMixer` for classification application. Args: config (`PatchTSMixerConfig`): Configuration. Returns: `None`. r-csdt|t||_t|d|_|j|_|jdvr$t|j |j d|_ nd|_ |j r0| dSdS)NrrMrLTr0r@) rrr)rrr9rrzInjectScalerStatistics4Dr0r@ inject_scaler"r3rr r!rs    z0PatchTSMixerForTimeSeriesClassification.__init__NFTr target_valuesrr;r#rAc Cstj}|dur |n|j}|j|||d}t|trt|}|jdur0|j|j |j |j d|_ | |j }|durD|durD|||} nd} |sWtdd| ||j |j fDSt| ||j |j dS)a past_values (`torch.FloatTensor` of shape `(batch_size, seq_length, num_input_channels)`): Context values of the time series. For a pretraining task, this denotes the input time series to predict the masked portion. For a forecasting task, this denotes the history/past time series values. Similarly, for classification or regression tasks, it denotes the appropriate context values of the time series. For univariate time series, `num_input_channels` dimension should be 1. For multivariate time series, it is greater than 1. target_values (`torch.FloatTensor` of shape `(batch_size, target_len, num_input_channels)` for forecasting, `(batch_size, num_targets)` for regression, or `(batch_size,)` for classification, *optional*): Target values of the time series, that serve as labels for the model. The `target_values` is what the Transformer needs during training to learn to output, given the `past_values`. Note that, this is NOT required for a pretraining task. For a forecasting task, the shape is be `(batch_size, target_len, num_input_channels)`. Even if we want to forecast only specific channels by setting the indices in `prediction_channel_indices` parameter, pass the target data with all channels, as channel Filtering for both prediction and target will be manually applied before the loss computation. For a classification task, it has a shape of `(batch_size,)`. For a regression task, it has a shape of `(batch_size, num_targets)`. return_loss (`bool`, *optional*): Whether to return the loss in the `forward` call. Nr0r^Tcsr$r"r r%r r r!r=r'zBPatchTSMixerForTimeSeriesClassification.forward..r?)r8rZCrossEntropyLossrrrrr(rorr r r9rrj) rrrprr;r#r3rAr`rBr r r!r$sB $     z/PatchTSMixerForTimeSeriesClassification.forwardrC)r%r&r'r(rrr r8r9rrrjr$r*r r rr!rks* rkc@r1)PatchTSMixerForRegressionOutputa Output type of [`PatchTSMixerForRegressionOutput`]. Args: regression_outputs (`torch.FloatTensor` of shape `(batch_size, num_targets)`): Prediction output from the regression head. last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, d_model)`): Backbone embeddings before passing through the head. hidden_states (`tuple(torch.FloatTensor)`, *optional*): Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`): Total loss. Nr3regression_outputsrr) r%r&r'r(r3rr8rrrrrrrr r r r!rqOr6rqcsHeZdZd dededeffdd Zdejdejd ejfd d ZZS) rnr4r0r@ expansioncs`tt|d|||_t||||_tdd||_td|d|_||_dSr) rrrrinverse_trans_expansioninverse_trans_compressionmap_scale_expansionmap_scale_compressionr@)rr0r@rsrr r!rfs  z!InjectScalerStatistics4D.__init__r#r r cCs|dd}|d}|dd|jd}|dd}|d}|dd|jd}tj||gdd}||}||}tj||gdd}||}| |}|S)a Args: inputs (`torch.Tensor` of shape `(batch_size, num_input_channels, num_patch, d_model)`) loc (`torch.Tensor` of shape `(batch_size, 1, num_input_channels)`) scale (`torch.Tensor` of shape `(batch_size, 1, num_input_channels)`) Returns: `torch.Tensor` of shape `(batch_size, num_input_channels, num_patch, d_model)` rrrr) r5rFrr@r8catrvrwrtru)rr#r r rLstdevZ concat_statsr r r!r$os       z InjectScalerStatistics4D.forward)r4) r%r&r'r)rr8r9r$r*r r rr!rnes$ rncszeZdZdZdeffdd Ze    ddejd e ejd e e d e d e e d e f ddZ dejd e fddZZS)PatchTSMixerForRegressionz `PatchTSMixer` for regression application. Args: config (`PatchTSMixerConfig`): Configuration. Returns: `None`. r-cst|t||_|j|_|j|_|j|_|j|_|jdkr$d|_n tt t d}| |j}|dur<||j d|_nt d|j|jdvrSt|j|jd|_nd|_t||jd|_|jrg|dSdS)NrRrSrrVrlrmrW)rrr)rr3rrrXr r r rYrrNrzrnr0r@rorr9r")rr-rZr[rr r!rs4      z"PatchTSMixerForRegression.__init__NFTrrprr;r#rAc sDjdkr tjdd}n jdkrt}ntd|dur|nj}j|||d}t|tr2t |}j durCj |j |j |j d|_ |j }|dur|d urjrjd krdt|d krdtd j|} tfd d|D}|| |} t| } n|||} nd} |stdd| ||j |jfDSt| ||j |jdS)a past_values (`torch.FloatTensor` of shape `(batch_size, seq_length, num_input_channels)`): Context values of the time series. For a pretraining task, this denotes the input time series to predict the masked portion. For a forecasting task, this denotes the history/past time series values. Similarly, for classification or regression tasks, it denotes the appropriate context values of the time series. For univariate time series, `num_input_channels` dimension should be 1. For multivariate time series, it is greater than 1. target_values (`torch.FloatTensor` of shape `(batch_size, target_len, num_input_channels)` for forecasting, `(batch_size, num_targets)` for regression, or `(batch_size,)` for classification, *optional*): Target values of the time series, that serve as labels for the model. The `target_values` is what the Transformer needs during training to learn to output, given the `past_values`. Note that, this is NOT required for a pretraining task. For a forecasting task, the shape is be `(batch_size, target_len, num_input_channels)`. Even if we want to forecast only specific channels by setting the indices in `prediction_channel_indices` parameter, pass the target data with all channels, as channel Filtering for both prediction and target will be manually applied before the loss computation. For a classification task, it has a shape of `(batch_size,)`. For a regression task, it has a shape of `(batch_size, num_targets)`. return_loss (`bool`, *optional*): Whether to return the loss in the `forward` call. rRrLr=rMr]Nr0r^TrUrzDtarget_values cannot be negative for negative_binomial distribution.csg|] }|djjqS)r)rr-r)ritemrr r!rsz5PatchTSMixerForRegression.forward..csr$r"r r%r r r!r r'z4PatchTSMixerForRegression.forward..)r3rrrr)r3rr@rMrNrrrrr(rorr r r9rr8any Exceptionr_rPrrq) rrrprr;r#r3rAr`r_rBr rr!r$sX #         z!PatchTSMixerForRegression.forwardcs^|j}||ddd}|j|jfddt|D}tj|ddd||jj }t |d S) a  Generate sequences of sample predictions from a model with a probability distribution head. Args: 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 target values. Return: [`SamplePatchTSMixerRegressionOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of samples, num_targets)`. NF)rrprcrar rbrrdr r!r:sz6PatchTSMixerForRegression.generate..rrrre) rXrr_rrrr8rfrr-rrJ)rrrXrgrhr rdr!ris  z"PatchTSMixerForRegression.generaterC)r%r&r'r(rrr r8r9rrrqr$rJrir*r r rr!rzs4 ' \rz)rr)r8rQrkrz)NFr)Nr)NN)Mr(rH dataclassesrtypingrrrr8Ztorch.nnrZtransformers.modeling_utilsrZtransformers.utilsrZtime_series_utilsr r r utilsr r Zutils.deprecationrZconfiguration_patchtsmixerrZ get_loggerr%r{Modulerr,r:rSr[rerprrrrrrrrr9rlistrr)rrrrrrrrrr(r)r2r8rDrFrJ distributions DistributionrMrPrQrjrkrqrnrz__all__r r r r!s       '11E-&)8G" >  E1=$7Ea_ "Wn( 7