o ۷iL@sdZddlZddlmZddlmZmZmZddlZddl m Z ddl m Z ddl mZddlmZdd lmZdd lmZdd lmZmZmZdd lmZmZd dlmZeeZ Gddde j!Z"Gddde j!Z#Gddde j!Z$Gddde j!Z%Gddde j!Z&Gddde j!Z'   dde j!dej(dej(dej(d eej(d!ee)d"e)d#eej(fd$d%Z*Gd&d'd'e j!Z+Gd(d)d)e j!Z,Gd*d+d+e j!Z-Gd,d-d-e j!Z.Gd.d/d/e j!Z/Gd0d1d1e j!Z0Gd2d3d3e j!Z1eGd4d5d5e Z2Gd6d7d7e j!Z3  8 dd9ej(d:e)d;ee4dd?Z7  dd9ej(d@ee4e6fd;ee4d=e6fdAdBZ8GdCdDdDe j!Z9GdEdFdFe j!Z:GdGdHdHe j!Z;GdIdJdJe j!ZGdQdRdRe2Z?eedSdNGdTdUdUeZ@edVdNGdWdXdXe2ZAeedYdNGdZd[d[eZBed\dNGd]d^d^e2ZCeed_dNGd`dadaeZDeedbdNGdcddddeZEeedbdNGdedfdfeZFdgejGjHdhej(diej(fdjdkZIddlej(dmeej(diej(fdndoZJGdpdqdqe2ZKeedrdNGdsdtdteZLGdudvdve2ZMeedwdNGdxdydyeZNGdzd{d{e j!ZOed|dNGd}d~d~e2ZPgdZQdS)zPyTorch PatchTSMixer model.N) dataclass)CallableOptionalUnion)PreTrainedModel) ModelOutput)FlashAttentionKwargs)ALL_ATTENTION_FUNCTIONS)Unpack)NegativeBinomialOutput NormalOutputStudentTOutput)auto_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__l/home/ubuntu/vllm_env/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_maskscalingrj head_maskc Ks|dur |dd}t||dd|} |dur| |} tjj| dd} |dur5| |dddd} tjj| ||j d} t| |} | dd } | | fS)Nrr9rrr)ptraining) sizer=matmulr:rrqsoftmaxviewrjr contiguous) rrrrrrrjrkwargs attn_weights attn_outputr#r#r$eager_attention_forwards  rcseZdZdZ     ddededed ed ed ed eeffd d Z    dde j dee j dee j dee j deede e dee j ee j eee j ffddZZS)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_statesrlayer_head_maskoutput_attentionsrrFcKs|du}|jdd\}} |r|jdn| } || d|jf} || d|jf} ||j| dd} |r4|n|}||j| dd}||j| dd}t}|jj dkr\t |jj }||| |||f|j shdn|j |j ||d|\}}||| d}||}||dfS)z#Input shape: Batch x Time x ChannelNrrr9eagerr)rjrrr)rdrrrr:rrrr1_attn_implementationr rrjrrcrr)r rrrrrris_cross_attentionbsztgt_lensrc_len q_input_shapekv_input_shape query_statescurrent_states key_states value_statesattention_interfacerrr#r#r$r(2s:      zPatchTSMixerAttention.forward)rFTFN)NNNF)r)r*r+r,r-floatboolrrrr=r>r r tupler(r.r#r#r!r$rsR"  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$rps(  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$rhsrcr/)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$r s "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$r2s    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$ Usz8PatchTSMixerForPredictionHead.forward..rrNc3s|] }|djfVqdS).N)rrr r#r$r[s.)rrr isinstancerr:rr hidden_featuresforecastr#rr$r(Ds      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$rjs&   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$rbsrc@s*eZdZUeed<dZdZdZddZdS)PatchTSMixerPreTrainedModelr1model 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 weightsrGrg?)rUrV?N)rr?r1rLrinitnormal_rBrbr4rdatazero_weightfill_r0r7rinit_std)r rr#r#r$ _init_weightss      z)PatchTSMixerPreTrainedModel._init_weightsN) r)r*r+r__annotations__base_model_prefixmain_input_namesupports_gradient_checkpointingrr#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)`. )rr rr#r#r$r(s z PatchTSMixerPretrainHead.forwardrr#r#r!r$r sr Fr& 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$r7sz$forecast_masking..rrznum_forecast_mask_patches z6 should be greater than 0 and less than total patches.cSs|dSNr9r#)xr#r#r$Isz"forecast_masking..)rr9rrNr)rr-rdr=rDrsumziprWrsortedrandpermrOrrr)r&r&rrforecast_mask_ratiosrrrrr!t_list total_length total_ratior ratiotemp_lenbatch1 patch_lenrbatch2permr$r#r#r$forecast_maskingsB    r8cr/)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_lengthrr  patch_striderWrrEsequence_start)r r1new_sequence_lengthr!r#r$rjs   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) rdrrWr=unfoldr r<r:r)r rrr;r#r#r$r({s zPatchTSMixerPatchify.forwardr<r#r#r!r$r9bsr9cr/)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 .) rDr%rCrrrr8r&rWr)r rZ masked_inputr!r#r#r$r(s$   zPatchTSMixerMasking.forwardr<r#r#r!r$rBs  rBc 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㈵>)rrhasattrrIrrJrKr8r!r#r$rs  zPatchTSMixerStdScaler.__init__robserved_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)`) rJrr9)r*rrJ clamp_minr=sqrtrK)r rrM denominatorlocvariancescaler#r#r$r(s  "zPatchTSMixerStdScaler.forward r)r*r+r,rrr=r>rr(r.r#r#r!r$rHsrHc rG) 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)NrIrrJTrK绽|= default_scale)rrrLrIrrJrKrYr8r!r#r$rs  zPatchTSMixerMeanScaler.__init__rrMrFc 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)rNTrOrminNrr) absr*rr=clamprYsqueeze ones_likewhererKrJ zeros_like) r rrMts_sum num_observedrU batch_sumbatch_observationsrY scaled_datar#r#r$r(s  zPatchTSMixerMeanScaler.forwardrVr#r#r!r$rWsrWc 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. r1cs@tt|dr |jnd|_t|dr|j|_dSd|_dS)NrIrrJT)rrrLrIrrJr8r!r#r$r0s  zPatchTSMixerNOPScaler.__init__NrrMrFcCsBtj|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)rrJ)r=r_rUrrJra)r rrMrUrSr#r#r$r(5s  zPatchTSMixerNOPScaler.forwardr%) r)r*r+r,rrr=r>rrr(r.r#r#r!r$rg+srgzS Base class for `PatchTSMixerEncoderOutput`, with potential hidden states. ) custom_introc@s:eZdZUdZdZeejed<dZ ee ejed<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,rjrr= FloatTensorrrrr#r#r#r$riFs ric 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. r1csbt||j|_t|j|j|_|jrt |d|_ nd|_ t |d|_ |j r/| dSdS)Nr)rruse_return_dictrrr r5patcherr@r?positional_encoderrmlp_mixer_encoder post_initr8r!r#r$ras   zPatchTSMixerEncoder.__init__FNrr return_dictrFcCsh|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..)rjr)rmrnrorprri)r rrrrpatchesrjrr#r#r$r(qs    zPatchTSMixerEncoder.forward)FN)r)r*r+r,rrrr=r>rrrrrir(r.r#r#r!r$rlXs rlzG Base class for model's outputs, with potential hidden states. c@seZdZUdZdZeejed<dZ ee ejed<dZ eejed<dZ eejed<dZ eejed<dZeejed<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. NrjrrZr!rSrU)r)r*r+r,rjrr=rkrrrrZr!rSrUr#r#r#r$rxs rxz= The PatchTSMixer Model for time-series forecasting. c sbeZdZddedeffdd Ze   ddejde ejd e ed e ed e f d d Z Z S)PatchTSMixerModelFr1 mask_inputcst||j|_t||_t||_|durt||_nd|_|j dkr,t ||_ n|j dks6|j dur.)rjrrZr!rSrU) rmr=r_r~r|r}r{rrrirjrrx) r rrrrrr!scaled_past_valuesrSrU patched_x enc_inputencoder_outputr#r#r$r(sD     zPatchTSMixerModel.forwardr)NFN) r)r*r+rrrrr=r>rrxr(r.r#r#r!r$rys"ryz> Output type of [`PatchTSMixerForPreTrainingOutput`]. c@^eZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<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_outputsrjr r)r*r+r,rrr=rkrrrjrrr#r#r#r$r#  rz. `PatchTSMixer` for mask pretraining. csbeZdZdeffdd Ze    ddejdeejd ee d e d ee d e f d dZ Z S)PatchTSMixerForPretrainingr1csLt|t|dd|_t|d|_|j|_|j|_|jr$|dSdS)NT)rzr) rrryrr head masked_lossrmrqr8r!r#r$rAs   z#PatchTSMixerForPretraining.__init__NFTrrr return_lossrrrFc 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. NTnone reductionrUrrrrrrrXcsrsr%r#rtr#r#r$rrvz5PatchTSMixerForPretraining.forward..rrrjr)rmrr=rMSELossrrrrxrrjrZrUr!r*rr) r rrrrrrr model_outputx_hatloss_valr#r#r$r(Ls@   $ z"PatchTSMixerForPretraining.forwardNFTN) r)r*r+rrrr=r>rrrr(r.r#r#r!r$r;s( rz= Output type of [`PatchTSMixerForPredictionOutput`]. c@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dZ eejed<dZeejed<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 NrrrjrrSrU)r)r*r+r,rrr=rkrrrjrrrSrUr#r#r#r$rs rz Base class for time series model's predictions outputs that contains the sampled values from the chosen distribution. c@$eZdZUdZdZeejed<dS)"SamplePatchTSMixerPredictionOutput 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,rrr=rkrr#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`. NrrrrZ)r=r`rar]r*rU)rrrweighted_tensor sum_weightsr#r#r$weighted_averages "  rcseZdZdZdeffdd Ze     ddejd e ejd e ejd e e d e d e e de fddZ e ddejd e ejdefddZZS)PatchTSMixerForPredictionz `PatchTSMixer` for forecasting application. Args: config (`PatchTSMixerConfig`): Configuration. Returns: `None`. r1cst||j|_|j|_|j|_|j|_|jdkrd|_n#|j}tt t d}| |jd}|dur:||d|_nt d|jt ||_t||jd|_|jrX|dSdS)Nmse student_tnormalnegative_binomialrUnknown distribution output r1r)rrrrmrnum_parallel_samplesrrrr r getrWryrrrrq)r r1rdistribution_output_map output_classr!r#r$rs0    z"PatchTSMixerForPrediction.__init__NFTrr future_valuesrrrrrFcCs|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.rSrUTcsrsr%r#rtr#r#r$rrvz4PatchTSMixerForPrediction.forward..)rrrjrrSrU)rrrrrWrmrrrrxrrjrr distributionrSrUrrr)r rrrrrrrrry_hatrrrSrUr#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) rrrrrSrUrr=stackr)r rrroutputssamplesr#rr$generates  z"PatchTSMixerForPrediction.generate)NNFTNr%)r)r*r+r,rrrr=r>rrrr(no_gradrrr.r#r#r!r$rsB   yrzK 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. Nrrrjrrr#r#r#r$rrrcsfeZdZdZdeffdd Ze    ddejd e ejd e e d e d e e d 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) rrryrrrrmrInjectScalerStatistics4Dr5rE inject_scalerqr8r!r#r$rs    z0PatchTSMixerForTimeSeriesClassification.__init__NFTr target_valuesrrrrrFc 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. NrrTcsrsr%r#rtr#r#r$r>rvzBPatchTSMixerForTimeSeriesClassification.forward..r)r=rCrossEntropyLossrmrrrrxrrjrSrUrrr) r rrrrrrrrrrr#r#r$r(sB $     z/PatchTSMixerForTimeSeriesClassification.forwardr)r)r*r+r,rrrr=r>rrrr(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_outputsrjr) r)r*r+r,rrr=rkrrrjrrr#r#r#r$rPrrcsHeZdZd 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$ris  z!InjectScalerStatistics4D.__init__r&rSrUcCs|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&rSrUrUstdev concat_statsr#r#r$r(rs       z InjectScalerStatistics4D.forward)r9) r)r*r+r-rr=r>r(r.r#r#r!r$rhs$ rz4 `PatchTSMixer` for regression application. cs~eZdZdeffdd Ze    ddejdeejd ee d e d ee d e f d dZ e dejd e fddZZS)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)rrryrrrrmrrr r rrrWrrr5rErrrrq)r r1rrr!r#r$rs4      z"PatchTSMixerForRegression.__init__NFTrrrrrrrFc 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..csrsr%r#rtr#r#r$r rv)rrrjr)rrrrrWrmrrrrxrrjrSrUrrr=any Exceptionrrrr) r rrrrrrrrrrrr#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$r8sz6PatchTSMixerForRegression.generate..rrrr) rrrrrr=rrr1rr)r rrrrr#rr$rs  z"PatchTSMixerForRegression.generater)r)r*r+rrrr=r>rrrr(rrrr.r#r#r!r$rs4' \r)rryrrrr)NrN)NFr)Nr)NN)Rr,rQ dataclassesrtypingrrrr=torch.nnrtransformers.modeling_utilsrtransformers.utilsrmodeling_flash_attention_utilsr modeling_utilsr processing_utilsr time_series_utilsr r rutilsrrconfiguration_patchtsmixerr get_loggerr)loggerModulerr0r?r]rfrsr>rrrrrrrrrrr listrr-r%r8r9rBrHrWrgrirlrxryrrrrr distributions Distributionrrrrrrrr__all__r#r#r#r$s         '17 XF-&)8G" >  E1=$7 EaT  "Xn( -