o eiL@sdZddlZddlmZddlmZddlZddlmZddl m Z ddl m Z ddl mZdd lmZdd lmZdd lmZdd lmZmZmZdd lmZmZmZddlmZe e!Z"Gdddej#Z$Gdddej#Z%Gdddej#Z&Gdddej#Z'Gdddej#Z(Gdddej#Z)  ddej#dej*dej*d ej*d!ej*dBd"e+dBd#e+d$eefd%d&Z,Gd'd(d(ej#Z-Gd)d*d*ej#Z.Gd+d,d,ej#Z/Gd-d.d.ej#Z0Gd/d0d0ej#Z1Gd1d2d2ej#Z2Gd3d4d4ej#Z3eGd5d6d6e Z4Gd7d8d8ej#Z5  9 dd:ej*d;e+de8f d?d@Z9  dd:ej*dAe6e8Bde8fdBdCZ:GdDdEdEej#Z;GdFdGdGej#ZGdLdMdMej#Z?eedNdOGdPdQdQe Z@GdRdSdSe4ZAeedTdOGdUdVdVe ZBedWdOGdXdYdYe4ZCeedZdOGd[d\d\e ZDed]dOGd^d_d_e4ZEeed`dOGdadbdbe ZFeedcdOGdddedee ZGeedcdOGdfdgdge ZHdhejIjJdiej*djej*fdkdlZKddmej*dnej*dBdjej*fdodpZLGdqdrdre4ZMeedsdOGdtdudue ZNGdvdwdwe4ZOeedxdOGdydzdze ZPGd{d|d|ej#ZQed}dOGd~dde4ZRgdZSdS)zPyTorch PatchTSMixer model.N)Callable) dataclass)PreTrainedModel) ModelOutput)initialization)FlashAttentionKwargs)ALL_ATTENTION_FUNCTIONS)Unpack)NegativeBinomialOutput NormalOutputStudentTOutput)TransformersKwargsauto_docstringlogging)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_layerSoftmax attn_softmax)selfrr __class__t/home/ubuntu/transcripts/venv/lib/python3.10/site-packages/transformers/models/patchtsmixer/modeling_patchtsmixer.pyr/s z#PatchTSMixerGatedAttention.__init__cCs|||}||}|SN)rr)r inputs attn_weightr#r#r$forward4sz"PatchTSMixerGatedAttention.forward)__name__ __module__ __qualname____doc__intrr( __classcell__r#r#r!r$r&src6eZdZdZdeffdd ZdejfddZZ S)PatchTSMixerBatchNormzP Compute batch normalization over the sequence length (time) dimension. configcs"ttj|j|jd|_dS)Neps)rrr BatchNorm1dd_modelnorm_eps batchnormr r1r!r#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) transposer7)r r&outputr#r#r$r(Ds   zPatchTSMixerBatchNorm.forward r)r*r+r,rrtorchTensorr(r.r#r#r!r$r0;sr0csNeZdZdZdeffdd ZededejfddZ de j fd d Z Z S) PatchTSMixerPositionalEncodingz' Class for positional encoding r1cs<t|jr|||_dStt|j |j |_dSr%) rruse_positional_encoding_init_pe position_encr Parameterr=zeros num_patchesr5r8r!r#r$rVs 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 ) NrandomT requires_gradsincosrrr9g@ FzN is not a valid positional encoder. Available types are 'random' and 'sincos'.)positional_encoding_typerrCr=randnrEr5rDarange unsqueezeexpmathlogsincosmeanstd ValueError)r1rBpositiondiv_termr#r#r$rA^s  (    z'PatchTSMixerPositionalEncoding._init_pe patch_inputcCs||j}|Sr%)rB)r rZ hidden_stater#r#r$r(rs z&PatchTSMixerPositionalEncoding.forward)r)r*r+r,rr staticmethodrrCrAr=r>r(r.r#r#r!r$r?Qs r?cr/)PatchTSMixerNormLayerzeNormalization block Args: config (`PatchTSMixerConfig`): Configuration. r1csFt|j|_d|jvrt||_dStj|j|j d|_dS)Nbatchr2) rrnorm_mlplowerr0normr LayerNormr5r6r8r!r#r$rs 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))` r^rrr9r)r_r`r=reshapeshapera)r r&inputs_reshapedr#r#r$r(s  zPatchTSMixerNormLayer.forwardr<r#r#r!r$r]xs r]cs,eZdZfddZdejfddZZS)PatchTSMixerMLPcsPt||j}t|||_t|j|_t|||_ t|j|_ dSr%) rrexpansion_factorrrfc1Dropoutdropoutdropout1fc2dropout2)r in_features out_featuresr1 num_hiddenr!r#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` )rkr functionalgelurhrlrm)r r&r#r#r$r(s  zPatchTSMixerMLP.forward)r)r*r+rr=r>r(r.r#r#r!r$rfs rfcr/)$PatchTSMixerChannelFeatureMixerBlockzThis module mixes the features in the channel dimension. Args: config (`PatchTSMixerConfig`): Configuration. r1cPtt||_|j|_t|j|j|d|_|jr&t|j|jd|_ dSdSNrnror1rr) rrr]ra gated_attnrfnum_input_channelsmlpr gating_blockr8r!r#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` rrr9r)rapermuterxr{rz)r r&residualoutr#r#r$r(s   z,PatchTSMixerChannelFeatureMixerBlock.forwardr<r#r#r!r$rssrsmodulequerykeyvalueattention_maskscalingrjkwargsc Ks|dur |dd}t||dd|}|dur||}tjj|dd}tjj|||jd}t||} | dd } | |fS)Nrr9rr)ptrainingr) sizer=matmulr:rrqsoftmaxrjr contiguous) rrrrrrrjr attn_weights attn_outputr#r#r$eager_attention_forwards  rcseZdZdZ     ddededed ed ed ed edBffd d Z   dde j de j dBde j dBdedBde e de e j e j dBe e j dBff ddZZS)PatchTSMixerAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrFTN embed_dim num_headsrj is_decoderbias is_causalr1cst||_||_||_|||_||_|j||jkr*td|jd|d|jd|_||_ ||_ t j |||d|_ t j |||d|_t j |||d|_t j |||d|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: ).r)r)rrrrrjhead_dimr1rWrrrrrk_projv_projq_projout_proj)r rrrjrrrr1r!r#r$rs&    zPatchTSMixerAttention.__init__ hidden_stateskey_value_statesroutput_attentionsrrFcKs|du}|jdd\}}|r|jdn|} ||d|jf} || d|jf} ||j| dd} |r4|n|} || j| dd}|| j| dd}t|j j t }||| |||f|j sbdn|j |j|d|\}}|||d}||}||dfS)z#Input shape: Batch x Time x ChannelNrrr9r)rjrr)rdrrviewr:rrr get_interfacer1_attn_implementationrrrjrrcrr)r rrrrris_cross_attentionbsztgt_lensrc_len q_input_shapekv_input_shape query_statescurrent_states key_states value_statesattention_interfacerrr#r#r$r(0s8    zPatchTSMixerAttention.forward)rFTFN)NNF)r)r*r+r,r-floatboolrrr=r>r rtupler(r.r#r#r!r$rsL" rc.eZdZdZdeffdd ZddZZS)PatchMixerBlockzxThis module mixes the patch dimension. Args: config (`PatchTSMixerConfig`): Configuration. r1cstt||_|j|_|j|_t|j|j|d|_|jr(t |j|jd|_ |jr>t |j |j |j|d|_t||_dSdS)Nrvrw)rrrjr1)rrr]ra self_attnrxrfrErzrr{rr5self_attn_headsrjself_attn_layer norm_attnr8r!r#r$rls(  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)rr9r) rarrdrcrr:rzrxr{r) r r[r} batch_sizen_varsrEr5hidden_state_reshapedx_attn_r~r#r#r$r(s      zPatchMixerBlock.forwardr)r*r+r,rrr(r.r#r#r!r$rdsrcr/)FeatureMixerBlockzThis module mixes the hidden feature dimension. Args: config (`PatchTSMixerConfig`): Configuration. r1crtru) rrr]rarxrfr5rzrr{r8r!r#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. )rarzrxr{)r rr}r~r#r#r$r(s   zFeatureMixerBlock.forwardr<r#r#r!r$rsrcr/)PatchTSMixerLayerz The `PatchTSMixer` layer that does all three kinds of mixing. Args: config (`PatchTSMixerConfig`): Configuration. r1csHtt|d|_t|d|_|j|_|jdkr"t|d|_dSdS)Nr1 mix_channel) rrr patch_mixerr feature_mixermoderschannel_feature_mixerr8r!r#r$rs    zPatchTSMixerLayer.__init__rcCs,|jdkr ||}||}||}|S)rr)rrrr)r rr#r#r$r(s   zPatchTSMixerLayer.forwardr<r#r#r!r$rs  rcs6eZdZdZdeffdd Zd defddZZS) PatchTSMixerBlockzThe main computing framework of the `PatchTSMixer` model. Args: config (`PatchTSMixerConfig`): Configuration. r1cs2tj}tfddt|D|_dS)Ncsg|]}tdqS)r)r.0rrr#r$ sz.PatchTSMixerBlock.__init__..)rr num_layersr ModuleListrangemixers)r r1rr!rr$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)r r[rall_hidden_states embeddingmodr#r#r$r( s   zPatchTSMixerBlock.forwardF) r)r*r+r,rrrr(r.r#r#r!r$rsrc0eZdZdZddeffdd ZddZZS) PatchTSMixerForPredictionHeadzqPrediction Head for Forecasting Args: config (`PatchTSMixerConfig`): Configuration. Nr1cs|t|j|_|jdur|jt|j|_|dur+t|j |j |j |_ n | |j |j |_ tjdd|_dS)N start_dim)rrprediction_channel_indicessortrri head_dropout dropout_layerrrEr5prediction_lengthbase_forecast_blockget_parameter_projectionFlattenflatten)r r1distribution_outputr!r#r$r.s    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)r:rzr#r#r$ Qsz8PatchTSMixerForPredictionHead.forward..rrNc3s|] }|djfVqdS).N)rrr r#r$rWs.)rrr isinstancerr:rr hidden_featuresforecastr#rr$r(@s      z%PatchTSMixerForPredictionHead.forwardr%rr#r#r!r$r&srcr) PatchTSMixerLinearHeadzLinear head for Classification and Regression. Args: config (`PatchTSMixerConfig`): Configuration. Nr1cst|j|_|j|_|jdur|j}nd}||_|dur.t|j|j ||j |_ n | |j|j ||_ |jdurGtj dd|_ntj dd|_t|j|_dS)Nrrr)rrhead_aggregation output_rangerErrrr5ry num_targets projectionrrrrirrj)r r1r mul_factorr!r#r$rfs&   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)`. rruse_last).rmax_poolravg_poolNrr) r:rmaxvaluesrUrrjrrrr=sigmoid)r rr#r#r$r(s        &zPatchTSMixerLinearHead.forwardr%rr#r#r!r$r^src@s6eZdZUeed<dZdZdZdZe ddZ dS) PatchTSMixerPreTrainedModelr1model past_values)timeFcCst|tr|jjdkrtj|jddddSdSt|tjtj frKt |j t |j t|dddurIt |jt |jt |jdSdSt|tr`t |jj t |jj dSt|tjr~tj|j d|jjd|j durt |j dSdSdS)zInitialize weightsrGrg?)rUrV running_meanN)rr?r1rLinitnormal_rBrrbr4zeros_rones_weightgetattrr running_varnum_batches_trackedr0r7rinit_std)r rr#r#r$ _init_weightss*         z)PatchTSMixerPreTrainedModel._init_weightsN) r)r*r+r__annotations__base_model_prefixmain_input_nameinput_modalitiessupports_gradient_checkpointingr=no_gradrr#r#r#r$rs rcr)PatchTSMixerPretrainHeadzcPretraining head. Args: config (`PatchTSMixerConfig`): Configuration. r1cs.tt|j|_t|j|j|_ dSr%) rrrrirrrr5 patch_length base_pt_blockr8r!r#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#r!r$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) rWrdrr-r=randrepeatonesargsortgatherrO masked_fillr)r&rrrrr num_channelssequence_length num_featuresrlen_keepnoisemask ids_shuffle ids_restore inputs_maskr#r#r$random_maskings& r(num_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$r9sz$forecast_masking..rrznum_forecast_mask_patches z6 should be greater than 0 and less than total patches.cSs|dSNr9r#)xr#r#r$Ksz"forecast_masking..)rr9rrNr)rr-rdr=rDrsumziprWrsortedrandpermrOrrr)r&r)rrforecast_mask_ratiosrrr r!r$t_list total_length total_ratiorratiotemp_lenbatch1 patch_lenrbatch2permr'r#r#r$forecast_maskingsB    r;cr/)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)` r1cst|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) rrcontext_lengthr r patch_striderWrrEsequence_start)r r1new_sequence_lengthr!r#r$rls   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 (rN) dimensionrstepr) rdr rWr@unfoldrr?r:r)r rr r;r#r#r$r(}s zPatchTSMixerPatchify.forwardr<r#r#r!r$r<dsr<cr/)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 r1csXt|j|_|j|_|j|_|j|_|j|_|j|_|jdur*t|j|_dSdSr%) rrrandom_mask_ratior mask_typer)rrr/r8r!r#r$rs  zPatchTSMixerMasking.__init__rZcCsr|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 rG)r&rrrrr)r&r)rrzInvalid mask type .) rGr(rFrrrr;r)rWr)r rZ masked_inputr$r#r#r$r(s$   zPatchTSMixerMasking.forwardr<r#r#r!r$rEs  rEc 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. r1csVtt|dr |jnd|_t|dr|jnd|_t|dr&|j|_dSd|_dS)N scaling_dimrkeepdimT minimum_scalegh㈵>)rrhasattrrLrrMrNr8r!r#r$rs  zPatchTSMixerStdScaler.__init__dataobserved_indicatorrFcCsz|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)`) rM?r9)r-rrM clamp_minr=sqrtrN)r rPrQ denominatorlocvariancescaler#r#r$r(s  "zPatchTSMixerStdScaler.forward r)r*r+r,rrr=r>rr(r.r#r#r!r$rKsrKc rJ) PatchTSMixerMeanScalerz Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data accordingly. r1csltt|dr |jnd|_t|dr|jnd|_t|dr#|jnd|_t|dr1|j|_dSd|_dS)NrLrrMTrN绽|= default_scale)rrrOrLrrMrNr^r8r!r#r$rs  zPatchTSMixerMeanScaler.__init__rPrQrFc 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)rRTrSrminNrr) absr-rr=clampr^squeeze ones_likewhererNrM zeros_like) r rPrQts_sum num_observedrZ batch_sumbatch_observationsr^ scaled_datar#r#r$r(s  zPatchTSMixerMeanScaler.forwardr[r#r#r!r$r\sr\c sXeZdZdZdeffdd Z d dejdejdBdeejejejffd 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. r1cs@tt|dr |jnd|_t|dr|j|_dSd|_dS)NrLrrMT)rrrOrLrrMr8r!r#r$r2s  zPatchTSMixerNOPScaler.__init__NrPrQrFcCsBtj|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)`) FrH)rrM)r=rdrUrrMrf)r rPrQrZrXr#r#r$r(7s  zPatchTSMixerNOPScaler.forwardr%r[r#r#r!r$rl-srlzS Base class for `PatchTSMixerEncoderOutput`, with potential hidden states. ) custom_introc@s:eZdZUdZdZejdBed<dZe ejdBed<dS)PatchTSMixerEncoderOutputa- 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,ror= FloatTensorrrrr#r#r#r$rnHs rnc sXeZdZdZdeffdd Ze  d dejde dBd e dBd e e Bfd d Z Z S)PatchTSMixerEncoderz Encoder for PatchTSMixer which inputs patched time-series and outputs patched embeddings. Args: config (`PatchTSMixerConfig`): Configuration. r1csbt||j|_t|j|j|_|jrt |d|_ nd|_ t |d|_ |j r/| dSdS)Nr)rruse_return_dictrrrr5patcherr@r?positional_encoderrmlp_mixer_encoder post_initr8r!r#r$rcs   zPatchTSMixerEncoder.__init__FNrr return_dictrFcKsh|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..)ror)rrrsrtrurrn)r rrrwrpatchesrorr#r#r$r(ss    zPatchTSMixerEncoder.forward)FN)r)r*r+r,rrrr=r>rrrnr(r.r#r#r!r$rqZsrqzG Base class for model's outputs, with potential hidden states. c@seZdZUdZdZejdBed<dZe ejdBed<dZ ejdBed<dZ ejdBed<dZ ejdBed<dZ ejdBed<dS) PatchTSMixerModelOutputa 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. NrorrZr$rXrZ)r)r*r+r,ror=rprrrrZr$rXrZr#r#r#r$r}s r}z= The PatchTSMixer Model for time-series forecasting. c sbeZdZddedeffdd Ze   ddejdejdBd edBd edBd e f d d Z Z S)PatchTSMixerModelFr1 mask_inputcst||j|_t||_t||_|durt||_nd|_|j dkr,t ||_ n|j dks6|j dur.)rorrZr$rXrZ) rrr=rdrrrrrrrnrorr}) r rrrrwrr$scaled_past_valuesrXrZ patched_x enc_inputencoder_outputr#r#r$r(sD     zPatchTSMixerModel.forwardr)NFN) r)r*r+rrrrr=r>r}r(r.r#r#r!r$r~s"r~z> Output type of [`PatchTSMixerForPreTrainingOutput`]. c@^eZdZUdZdZejdBed<dZejdBed<dZ ejdBed<dZ e ejdBed<dS) PatchTSMixerForPreTrainingOutputa@ loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`): Total loss prediction_outputs (`torch.FloatTensor` of shape `(batch_size, num_input_channels, num_patches, patch_length)`): Prediction output from the pretrain 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. Nlossprediction_outputsror r)r*r+r,rr=rprrrorrr#r#r#r$r'  rz. `PatchTSMixer` for mask pretraining. csbeZdZdeffdd Ze    ddejdejdBd edBd ed edBd e f d dZ Z S)PatchTSMixerForPretrainingr1csLt|t|dd|_t|d|_|j|_|j|_|jr$|dSdS)NT)rr) rrr~rrhead masked_lossrrrvr8r!r#r$rEs   z#PatchTSMixerForPretraining.__init__NFTrrr return_lossrwrFc Ks|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. NTnone reductionrUrrrwrrr]csrxr%r#ryr#r#r$rr{z5PatchTSMixerForPretraining.forward..rrror)rrrr=rMSELossrrrr}rrorZrUr$r-rr) r rrrrrwrr model_outputx_hatloss_valr#r#r$r(Ps@   $ z"PatchTSMixerForPretraining.forwardNFTN) r)r*r+rrrr=r>rrr(r.r#r#r!r$r?s( rz= Output type of [`PatchTSMixerForPredictionOutput`]. c@seZdZUdZdZejdBed<dZejdBed<dZ ejdBed<dZ e ejdBed<dZ ejdBed<dZ ejdBed<dS) PatchTSMixerForPredictionOutputaD loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`): Total loss. 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. 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 NrrrorrXrZ)r)r*r+r,rr=rprrrorrrXrZr#r#r#r$rs rz Base class for time series model's predictions outputs that contains the sampled values from the chosen distribution. c@$eZdZUdZdZejdBed<dS)"SamplePatchTSMixerPredictionOutput sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, prediction_length, number_channels)`): Sampled values from the chosen distribution. N sequencesr)r*r+r,rr=rprr#r#r#r$r rc@r)"SamplePatchTSMixerRegressionOutputrNrrr#r#r#r$rrrinputtargetrFcCs || S)zc Computes the negative log likelihood loss from input distribution with respect to target. )log_prob)rrr#r#r$nlls r 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`. NrrrTr_)r=rerfrbr-rU)rrrweighted_tensor sum_weightsr#r#r$weighted_averages "  rcseZdZdZdeffdd Ze     ddejd ejdBd ejdBd e dBd e d e dBde fddZ e ddejd ejdBde fddZZS)PatchTSMixerForPredictionz `PatchTSMixer` for forecasting application. Args: config (`PatchTSMixerConfig`): Configuration. Returns: `None`. r1cst||j|_|j|_|j|_|j|_|jdkrd|_n"|j}tt t d}| |j}|dur9||d|_nt d|jt ||_t||jd|_|jrW|dSdS)Nmse student_tnormalnegative_binomialrUnknown distribution output r1r)rrrrrrnum_parallel_samplesrrr r r getrWr~rrrrv)r r1rdistribution_output_map output_classr!r#r$rs0     z"PatchTSMixerForPrediction.__init__NFTrr future_valuesrrrwrFcKs|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. rrUrr2Invalid loss function: Allowed values: mse and nllNr.rXrZTcsrxr%r#ryr#r#r$rr{z4PatchTSMixerForPrediction.forward..)rrrorrXrZ)rrrrrWrrrrrr}rrorr distributionrXrZrrr)r rrrrrrwrrry_hatrrrXrZr#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)rrrrrcg|]}qSr#samplerrr#r$rsz6PatchTSMixerForPrediction.generate..rrr) rrrrrXrZrr=stackr)r rrroutputssamplesr#rr$generates  z"PatchTSMixerForPrediction.generate)NNFTNr%)r)r*r+r,rrrr=r>rrr(r rrr.r#r#r!r$rsB   zrzK Output type of [`PatchTSMixerForTimeSeriesClassificationOutput`]. c@r)-PatchTSMixerForTimeSeriesClassificationOutputaP loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`): Total loss. 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. Nrrrorrr#r#r#r$rrrcsfeZdZdZdeffdd Ze    ddejd ejdBd e dBd e d e dBd e f ddZ Z S)'PatchTSMixerForTimeSeriesClassificationz `PatchTSMixer` for classification application. Args: config (`PatchTSMixerConfig`): Configuration. Returns: `None`. r1csdt|t||_t|d|_|j|_|jdvr$t|j |j d|_ nd|_ |j r0| dSdS)NrrVrUTr5rE) rrr~rrrrrrInjectScalerStatistics4Dr5rE inject_scalervr8r!r#r$rs    z0PatchTSMixerForTimeSeriesClassification.__init__NFTr target_valuesrrrwrFc Kstj}|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. NrrTcsrxr%r#ryr#r#r$rEr{zBPatchTSMixerForTimeSeriesClassification.forward..r)r=rCrossEntropyLossrrrrrr}rrorXrZrrr) r rrrrrwrrrrrr#r#r$r(sB %     z/PatchTSMixerForTimeSeriesClassification.forwardr) r)r*r+r,rrrr=r>rrr(r.r#r#r!r$rs* rz= Output type of [`PatchTSMixerForRegressionOutput`]. c@r)PatchTSMixerForRegressionOutputaM loss (*optional*, returned when `y` is provided, `torch.FloatTensor` of shape `()`): Total loss. 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. Nrregression_outputsror) r)r*r+r,rr=rprrrorrr#r#r#r$rWrrcsHeZdZd dededeffdd Zdejdejd ejfd d ZZS) rr9r5rE expansioncs`tt|d|||_t||||_tdd||_td|d|_||_dSr*) rrrrinverse_trans_expansioninverse_trans_compressionmap_scale_expansionmap_scale_compressionrE)r r5rErr!r#r$rps  z!InjectScalerStatistics4D.__init__r&rXrZcCs|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) r:rOrrEr=catrrrr)r r&rXrZrUstdev concat_statsr#r#r$r(ys       z InjectScalerStatistics4D.forward)r9) r)r*r+r-rr=r>r(r.r#r#r!r$ros$ rz4 `PatchTSMixer` for regression application. cs~eZdZdeffdd Ze    ddejdejdBd edBd ed edBd e f d dZ e dejd e fddZ ZS)PatchTSMixerForRegressionr1cst|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)Nrrrrrrr)rrr~rrrrrrr r r rrrWrrr5rErrrrv)r r1rrr!r#r$rs4      z"PatchTSMixerForRegression.__init__NFTrrrrrwrFc  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. rrUrrrNrrTrrzDtarget_values cannot be negative for negative_binomial distribution.c3s |] }|djjVqdS)rN)rr1r)ritemrr#r$rsz4PatchTSMixerForRegression.forward..csrxr%r#ryr#r#r$rr{)rrror)rrrrrWrrrrrr}rrorXrZrrr=any Exceptionrrrr) r rrrrrwrrrrrrr#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)rrrcrr#rrrr#r$r@sz6PatchTSMixerForRegression.generate..rrrr) rrrrrr=rrr1rr)r rrrrr#rr$r"s  z"PatchTSMixerForRegression.generater)r)r*r+rrrr=r>rrr(r rrr.r#r#r!r$rs4' ]r)rr~rrrr)Nr)NFr)Nr)NN)Tr,rQcollections.abcr dataclassesrr=torch.nnrtransformers.modeling_utilsrtransformers.utilsrrrmodeling_flash_attention_utilsrmodeling_utilsr processing_utilsr time_series_utilsr r r utilsrrrconfiguration_patchtsmixerr get_loggerr)loggerModulerr0r?r]rfrsr>rrrrrrrrrrrlistrr-r(r;r<rErKr\rlrnrqr}r~rrrrr distributions Distributionrrrrrrrr__all__r#r#r#r$s            '17 VF-&)8G" > E1=$7 FbU  "Yo( .