o Zhs` @sxdZddlZddlmZddlmZmZmZddlZddlm Z ddl m Z ddl m Z dd lmZdd lmZmZmZdd lmZmZmZd d lmZeeZGddde jZGddde jZ   dadej de!dee"de#de$f ddZ%  dbdej dee"e$fdee"de$fddZ&Gddde jZ'Gdd d e jZ(Gd!d"d"e jZ)eGd#d$d$eZ*Gd%d&d&e jZ+Gd'd(d(e jZ,Gd)d*d*e*Z-eGd+d,d,eZ.eGd-d.d.eZ/eGd/d0d0eZ0eGd1d2d2eZ1eGd3d4d4eZ2eGd5d6d6eZ3d7ej4j5d8ej d9ej fd:d;Z6dcdd?Z7Gd@dAdAe jZ8GdBdCdCe jZ9GdDdEdEe jZ:GdFdGdGe jZ;eGdHdIdIe*ZGdPdQdQe jZ?edRdMGdSdTdTe*Z@edUdMGdVdWdWe jZAedXdMGdYdZdZe*ZBGd[d\d\e jZCed]dMGd^d_d_e*ZDgd`ZEdS)dzPyTorch PatchTST model.N) dataclass)OptionalTupleUnion)nn)ACT2CLS)BaseModelOutput)PreTrainedModel)NegativeBinomialOutput NormalOutputStudentTOutput) ModelOutputauto_docstringlogging)PatchTSTConfigcseZdZdZ     ddededed ed ed ed eeffd d Z de j dedefddZ     dde j dee j dee e j dee j dee j dede e j ee j ee e j ffddZZS)PatchTSTAttentionz=Multi-headed attention from 'Attention Is All You Need' paperFTN embed_dim num_headsdropout is_decoderbias is_causalconfigcst||_||_||_|||_||_|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`: ).g࿩r)super__init__rrrhead_dimr ValueErrorscalingrrrLineark_projv_projq_projout_proj)selfrrrrrrr __class__]/var/www/auris/lib/python3.10/site-packages/transformers/models/patchtst/modeling_patchtst.pyr's&    zPatchTSTAttention.__init__tensorseq_lenbszcCs ||||j|jddS)Nr)viewrr transpose contiguous)r(r-r.r/r+r+r,_shapeFs zPatchTSTAttention._shape hidden_stateskey_value_statespast_key_valueattention_masklayer_head_maskoutput_attentionsreturncCsr|du}|\}} } |||j} |r.|dur.|djd|jdkr.|d} |d} nZ|rE|||d|} |||d|} nC|durt|||d|} |||d|} tj|d| gdd} tj|d| gdd} n|||d|} |||d|} |j r| | f}||j d|j f}|| | |j |} | j |} | j |} | d}t| | dd}|||j | |fkrtd||j | |fd||dur ||d| |fkrtd |d| |fd|| ||j | ||}| ||j | |}tjj|dd}|durL||j fkr1td |j fd|| dddd| ||j | |}| ||j | |}|rc| ||j | |}| ||j | |}nd}tjj||j|jd }t|| }|||j | |j fkrtd ||j | |j fd|| ||j | |j }|dd}| || |j}||}|||fS) z#Input shape: Batch x Time x ChannelNrr0rdimz$Attention weights should be of size z , but is z!Attention mask should be of size z/Head mask for a single layer should be of size )ptrainingz `attn_output` should be of size )sizer&r"shaper4r$r%torchcatrrr r1reshapeZbmmr2r!rZ functionalZsoftmaxrr@rr')r(r5r6r7r8r9r:Zis_cross_attentionr/Ztgt_len_Z query_statesZ key_statesZ value_statesZ proj_shapeZsrc_len attn_weightsZattn_weights_reshapedZ attn_probs attn_outputr+r+r,forwardIs        "    zPatchTSTAttention.forward)rFTFN)NNNNF)__name__ __module__ __qualname____doc__intfloatboolrrrrCTensorr4rrI __classcell__r+r+r)r,r$sV rc6eZdZdZdeffdd ZdejfddZZ S)PatchTSTBatchNormzP Compute batch normalization over the sequence length (time) dimension. rcs"ttj|j|jd|_dS)NZeps)rrrZ BatchNorm1dd_modelnorm_eps batchnormr(rr)r+r,rs zPatchTSTBatchNorm.__init__inputscCs"|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)` rr0)r2rX)r(rZoutputr+r+r,rIs   zPatchTSTBatchNorm.forward rJrKrLrMrrrCrQrIrRr+r+r)r,rTsrTFrZ 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.deviceNr<r=)r>index.r) r!rBrbrNrCZrandrepeatZonesZargsortgather unsqueeze masked_fillrP)rZr]r^r_r` batch_size num_channelssequence_length num_featuresrbZlen_keepnoisemaskZ ids_shuffleZ ids_restore inputs_maskr+r+r,random_maskings& rpnum_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+.0rFr+r+r, ,sz$forecast_masking..rarznum_forecast_mask_patches z6 should be greater than 0 and less than total patches.cSs|dS)Nr0r+)xr+r+r,>sz"forecast_masking..)keyr0r<rNrd) isinstancerNrBrCzerosrbsumzipr!appendsortedZrandpermrgrerhrP)rZrqr^r`Zforecast_mask_ratiosrirjrkrlrnZt_list total_length total_ratio patch_lengthratioZtemp_lenZbatch1Z patch_lenrFZbatch2permror+r+r,forecast_maskingsB    rcrS)PatchTSTPatchifyz A class to patchify the time series sequence into different patches Returns: `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)` rcst|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_lengthrkr patch_strider!max num_patchessequence_start)r(rZnew_sequence_lengthr)r+r,r^s   zPatchTSTPatchify.__init__ past_valuescCsp|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)` zInput sequence length (z%) doesn't match model configuration (rN) dimensionrAstep) rBrkr!rZunfoldrrr2r3)r(rrkr[r+r+r,rIos zPatchTSTPatchify.forwardr\r+r+r)r,rVsrcrS)PatchTSTMaskinga Class to perform random or forecast masking. Parameters: config (`PatchTSTConfig`): model config Returns: x_mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`) Masked patched input mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`) Bool tensor indicating True on masked points rcsXt|j|_|j|_|j|_|j|_|j|_|j|_|jdur*t|j|_dSdSN) rrrandom_mask_ratior_ mask_typerqr^r`r}rYr)r+r,rs  zPatchTSTMasking.__init__ patch_inputcCsr|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 random)rZr]r^r_r`Zforecast)rZrqr^r`zInvalid mask type .) rrprr^r_r`rrqr!rP)r(rZ masked_inputrnr+r+r,rIs$   zPatchTSTMasking.forwardr\r+r+r)r,rs  rcs@eZdZdZdeffdd Zd dejdee fdd Z Z S) PatchTSTEncoderLayerz PatchTST encoder layer rc st|j|_t|j|j|jd|_|jdkrt |jnt |_ |j dkr/t||_n|j dkr?t j|j|jd|_nt|j d|jr}|jdkrUt |jnt |_|j dkret||_n|j dkrut j|j|jd|_nt|j dt t j|j|j|jdt|j|jdkrt |jnt t j|j|j|jd|_|jdkrt |jnt |_|j dkrt||_n|j dkrt j|j|jd|_nt|j d|j|_dS)N)rrrrrXZ layernormrUz$ is not a supported norm layer type.r) rrchannel_attentionrrVZnum_attention_headsZattention_dropout self_attnZ path_dropoutrDropoutIdentity dropout_path1Z norm_typerTnorm_sublayer1 LayerNormrWr! dropout_path2norm_sublayer2Z Sequentialr#Zffn_dimrrZactivation_functionZ ff_dropoutff dropout_path3norm_sublayer3pre_normrYr)r+r,rsB               zPatchTSTEncoderLayer.__init__N hidden_stater:c Cs|j\}}}}|||||}|jr(|j|||d\}}} |||}n|j||d\}}} ||||}|||||}|jr|dd }|||||}|jrp|j| ||d\}} } || |}n|j||d\}} } | || |}|||||}|dd }|||||}|jr|| | ||}n ||| | |}|||||}|f} |r| |jr|| fn|f7} | S)a Parameters: hidden_state (`torch.Tensor` of shape `(batch_size, num_channels, sequence_length, d_model)`, *required*): Past values of the time series output_attentions (`bool`, *optional*): Whether or not to return the output attention of all layers Return: `torch.Tensor` of shape `(batch_size, num_channels, sequence_length, d_model)` )r5r:r0r)rBr1rrrrrErr2r3rrrrr) r(rr:rinum_input_channelsrkrVrHrGrFZchannel_attn_weightsoutputsr+r+r,rIsF       zPatchTSTEncoderLayer.forwardr) rJrKrLrMrrrCrQrrPrIrRr+r+r)r,rs"1rc@s.eZdZeZdZdZdZddZd ddZ dS) PatchTSTPreTrainedModelmodelrFcCst|tr&|jjrtjj|jdd|jjdkr$tjj|j ddddSdSt|tj r;|j j |jj ddSt|trQ|jj j |jjj ddSt|tjtjfrr|jj jd|jjd|j durt|j j dSdSdS) z$ Initialize weights g{Gz?)stdrrg?)meanr?N)rxPatchTSTPositionalEncodingr use_cls_tokenrinitZnormal_ cls_tokenpositional_encoding_type position_encrrdataZzero_weightZfill_rTrXr#ZConv1dZinit_std)r(moduler+r+r, _init_weightsSs$      z%PatchTSTPreTrainedModel._init_weightscCst|tr ||_dSdSr)rxPatchTSTEncodergradient_checkpointing)r(rvaluer+r+r,_set_gradient_checkpointingis  z3PatchTSTPreTrainedModel._set_gradient_checkpointingN)F) rJrKrLrZ config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrrr+r+r+r,rLsrc2eZdZdeffdd ZdejfddZZS)PatchTSTEmbeddingrcslt|j|_|j|_|jrt|j|j|_dSt |_t |jD]}|j t|j|jq%dSr) rrrshare_embeddingrr#rrVinput_embedding ModuleListranger|)r(rrFr)r+r,ros  zPatchTSTEmbedding.__init__rcsjjd}|jkrtdjd|djr }|Sfddt|D}tj|dd}|S)a% Parameters: patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*): Patch input for embedding return: `torch.Tensor` of shape `(batch_size, num_channels, num_patches, d_model)` rz&The defined number of input channels (zQ) in the config has to be the same as the number of channels in the batch input (rc s2g|]}j|dd|ddddfqSr)rrsirr(r+r,rts2z-PatchTSTEmbedding.forward..r=)rBrr!rrrrCstack)r(rrZ embeddingsr+rr,rI{s   zPatchTSTEmbedding.forward rJrKrLrrrCrQrIrRr+r+r)r,rns rcsVeZdZdZdedeffdd Zedededej fddZ d e j fd d Z ZS) rz' Class for positional encoding rrcszt|j|_|j|_|jr!ttddd|j|_ |d7}| |||_ |j dkr6t |j |_ dSt|_ dS)Nrr)rrrrr ParameterrCryrVr_init_perpositional_dropoutrrr(rrr)r+r,rs z#PatchTSTPositionalEncoding.__init__r;cCs|jdkrtjt||jdd}|S|jdkrst||j}td|d}t td|jdt d|j }t |||dddddf<t |||dddddf<||}||d }tj|d d}|St|jd ) NrTZ requires_gradZsincosrrr0g@ FzN is not a valid positional encoder. Available types are 'random' and 'sincos'.)rrrrCZrandnrVryZarangergexpmathlogsincosrrr!)rrrpositionZdiv_termr+r+r,rs  (    z#PatchTSTPositionalEncoding._init_percCs|jr8|||jddddf}|j|jddddf}||jd|jdd}tj||fdd}|S|||j}|S)Nrrr<r0r=) rrrrexpandrBrrCrD)r(rrZ cls_tokensrr+r+r,rIs z"PatchTSTPositionalEncoding.forward)rJrKrLrMrrNr staticmethodrrrrCrQrIrRr+r+r)r,rs rc sTeZdZdZdedeffdd Z  d dejde e d e e d e fd d Z Z S)rz PatchTST Encoder rrcsTtd|_t|_t||_tfddt j D|_ | dS)NFcsg|]}tqSr+)rrrr+r,rtz,PatchTSTEncoder.__init__..) rrrrembedderrpositional_encoderrrrZnum_hidden_layerslayers post_initrr)rr,rs     zPatchTSTEncoder.__init__Nroutput_hidden_statesr:r;c Cs|dur|n|jj}|dur|n|jj}||}||}|r"dnd}|r(dnd}|jD]}|r6||f}|||d}|d}|rI||df}q-t|||dS)a Parameters: patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*): Past values of the time series output_hidden_states (bool, optional): Indicates if hidden states should be outputted. output_attentions (bool, optional): Indicates if attentions should be outputted. return: `BaseModelOutput` Nr+)rr:rr)last_hidden_stater5 attentions)rr:rrrrr ) r(rrr:rencoder_statesZall_attentionsZ encoder_layerZ layer_outputsr+r+r,rIs        zPatchTSTEncoder.forwardNN)rJrKrLrMrrNrrCrQrrPr rIrRr+r+r)r,rsrc@seZdZUdZdZeejed<dZ ee ejed<dZ ee ejed<dZ eejed<dZ eejed<dZeejed<dZeejed <dS) PatchTSTModelOutputa Base class for model's outputs, with potential hidden states. Parameters: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`): Sequence of hidden-states at the output of the last layer of the model. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, num_channels, height, width)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. mask: (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches)`, *optional*) Bool masked tensor indicating which patches are masked loc: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*) Mean of the input data (batch_size, sequence_length, num_channels) over the sequence_length scale: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*) Std of the input data (batch_size, sequence_length, num_channels) over the sequence_length patch_input (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`): Patched input to the Transformer Nrr5rrnlocscaler)rJrKrLrMrrrC FloatTensor__annotations__r5rrrnrrrr+r+r+r,rs rc@beZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)PatchTSTForPretrainingOutputa Output type of [`PatchTSTForPretraining`]. Parameters: loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): MSE loss. prediction_outputs (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction outputs of the time series modeling heads. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nlossprediction_outputr5r) rJrKrLrMrrrCrrrr5rrr+r+r+r,r' rc@r)PatchTSTForRegressionOutputa Output type of [`PatchTSTForRegression`]. Parameters: loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): MSE loss. regression_outputs (`torch.FloatTensor` of shape `(batch_size, num_targets)`): Regression outputs of the time series modeling heads. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nrregression_outputsr5r) rJrKrLrMrrrCrrrr5rrr+r+r+r,rDrrc@seZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dZ eejed<dZeejed<dS) PatchTSTForPredictionOutputaR Output type of [`PatchTSTForPrediction`]. Parameters: loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): MSE loss. prediction_outputs (`torch.FloatTensor` of shape `(batch_size, prediction_length, -1)`): Prediction outputs of the time series modeling heads. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. loc: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*) Mean of the input data (batch_size, sequence_length, num_channels) over the sequence_length scale: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*) Std of the input data (batch_size, sequence_length, num_channels) over the sequence_length Nrprediction_outputsr5rrr)rJrKrLrMrrrCrrrr5rrrrr+r+r+r,ras rc@r)PatchTSTForClassificationOutputaR Output type of [`PatchTSTForClassification`]. Parameters: loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss. prediction_logits (`torch.FloatTensor` of shape `(batch_size, num_targets)`): Prediction scores of the PatchTST modeling head (scores before SoftMax). hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nrprediction_logitsr5r) rJrKrLrMrrrCrrrr5rrr+r+r+r,rs rc@s$eZdZUdZdZeejed<dS)SamplePatchTSTOutputa! Base class for time series model's predictions outputs that contains the sampled values from the chosen distribution. Parameters: sequences `(batch_size, num_samples, prediction_length, num_targets)`): Sampled values from the chosen distribution. N sequences) rJrKrLrMrrrCrrr+r+r+r,rs  rinputtargetr;cCs || S)zc Computes the negative log likelihood loss from input distribution with respect to target. )Zlog_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`. Nrr=rmin)rCwhere zeros_likeclamprzr)rrr>Zweighted_tensorZ sum_weightsr+r+r,weighted_averages "  rc PeZdZdZdeffdd Zdejdejdeejejejffdd Z Z S) PatchTSTStdScalerz 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. rcsVtt|dr |jnd|_t|dr|jnd|_t|dr&|j|_dSd|_dS)N scaling_dimrkeepdimT minimum_scalegh㈵>)rrhasattrrr>rrrYr)r+r,rs  zPatchTSTStdScaler.__init__robserved_indicatorr;cCsz|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)`) rrr0)rzr>rZ clamp_minrCsqrtr)r(rr denominatorrZvariancerr+r+r,rIs  "zPatchTSTStdScaler.forward rJrKrLrMrrrCrQrrIrRr+r+r)r,rsrc r) PatchTSTMeanScalerz Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data accordingly. rcsltt|dr |jnd|_t|dr|jnd|_t|dr#|jnd|_t|dr1|j|_dSd|_dS)NrrrTr绽|= default_scale)rrrrr>rrr rYr)r+r,rs  zPatchTSTMeanScaler.__init__rrr;c 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)rTrrrNrr=) absrzr>rCrr Zsqueeze ones_likerrrr) r(rrZts_sumZ num_observedrZ batch_sumZbatch_observationsr Z scaled_datar+r+r,rIs  zPatchTSTMeanScaler.forwardrr+r+r)r,rsrc sXeZdZdZdeffdd Z d dejdeejde ejejejffd d Z Z S) PatchTSTNOPScalerz| Assigns a scaling factor equal to 1 along the first dimension, and therefore applies no scaling to the input data. rcs@tt|dr |jnd|_t|dr|j|_dSd|_dS)NrrrT)rrrrr>rrYr)r+r,r2s  zPatchTSTNOPScaler.__init__Nrrr;cCsBtj|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)`) Fr)r>r)rCr rr>rr)r(rrrrr+r+r,rI7s  zPatchTSTNOPScaler.forwardr) rJrKrLrMrrrCrQrrrIrRr+r+r)r,r -sr c sLeZdZdeffdd ZdejdejdeejejejffddZZ S) PatchTSTScalerrcsRt|jdks|jdurt||_dS|jdkr"t||_dSt||_dS)NrTr)rrr"rscalerrr rYr)r+r,rIs  zPatchTSTScaler.__init__rrr;cCs|||\}}}|||fS)a> Parameters: data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`): Input for scaler calculation observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`): Calculating the scale on the observed indicator. Returns: tuple of `torch.Tensor` of shapes (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`, `(batch_size, 1, um_input_channels)`) )r)r(rrrrr+r+r,rIRs zPatchTSTScaler.forward) rJrKrLrrrCrQrrIrRr+r+r)r,r Hs r csveZdZdeffdd Z     ddejdeejdeejdeed eed eed e e e ffd d Z Z S) PatchTSTModelrcsft|t||_t||_|j|_|jj}|jr!t||_ nt |_ t ||d|_ |dS)N)r)rrr rr patchifier do_mask_inputrrmaskingrrrencoderrrr)r+r,rfs      zPatchTSTModel.__init__Nrpast_observed_mask future_valuesrr: return_dictr;c Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur't|}|||\}}} ||} |jr@| | \} } n| | d} } |j | ||d} |sk| j | j | j f}|| || | f}tdd|DSt| j | j | j | || | dS)a Parameters: past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). future_values (`torch.BoolTensor` of shape `(batch_size, prediction_length, num_input_channels)`, *optional*): Future target values associated with the `past_values` output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers output_attentions (`bool`, *optional*): Whether or not to return the output attention of all layers return_dict (`bool`, *optional*): Whether or not to return a `ModelOutput` instead of a plain tuple. Returns: `PatchTSTModelOutput` or tuple of `torch.Tensor` (if `return_dict`=False or `config.return_dict`=False) Examples: ```python >>> from huggingface_hub import hf_hub_download >>> import torch >>> from transformers import PatchTSTModel >>> file = hf_hub_download( ... repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset" ... ) >>> batch = torch.load(file) >>> model = PatchTSTModel.from_pretrained("namctin/patchtst_etth1_pretrain") >>> # during training, one provides both past and future values >>> outputs = model( ... past_values=batch["past_values"], ... future_values=batch["future_values"], ... ) >>> last_hidden_state = outputs.last_hidden_state ```N)rrr:css|] }|dur|VqdSrr+)rsvr+r+r, z(PatchTSTModel.forward..)rr5rrnrrr)ruse_return_dictr:rrCr rrrrrrr5rtupler)r(rrrrr:rZscaled_past_valuesrrZpatched_valuesZ masked_valuesrnZencoder_outputrr+r+r,rIxs66  zPatchTSTModel.forwardNNNNN)rJrKrLrrrCrQrrPrrrrIrRr+r+r)r,rds, rcs<eZdZdZdeffdd ZdejdejfddZZ S) PatchTSTMaskPretrainHeadz- Pretraining head for mask modelling rcsHt|jdkrt|jnt|_t|j|j |_ |j |_ dSNr) rr head_dropoutrrrrr#rVrlinearrrYr)r+r,rs   z!PatchTSTMaskPretrainHead.__init__ embeddingr;cCs:|||}|jr|ddddddddf}|S)a Parameters: embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*): Embedding from the model Returns: `torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True Nr)r rr)r(r!r+r+r,rIs  z PatchTSTMaskPretrainHead.forwardr\r+r+r)r,rsrz* The PatchTST for pretrain model. )Z custom_introcsjeZdZdeffdd Z    d dejdeejdeedeed eed e e e ff d d Z Z S)PatchTSTForPretrainingrcs4t|d|_t|d|_t||_|dS)NTr)rrrrrrheadrrYr)r+r,rs    zPatchTSTForPretraining.__init__Nrrrr:rr;c Cs|dur|n|jj}|j||||dd}||j}tjdd}|||j} | jdd|j  |j d} |j } |sU|f|d d } | durQ| f| } | S| } | St | || |j d S) a Parameters: past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers output_attentions (`bool`, *optional*): Whether or not to return the output attention of all layers return_dict (`bool`, *optional*): Whether or not to return a `ModelOutput` instead of a plain tuple. Returns: `PatchTSTForPretrainingOutput` or tuple of `torch.Tensor` (if `return_dict`=False or `config.return_dict`=False) Examples: ```python >>> from huggingface_hub import hf_hub_download >>> import torch >>> from transformers import PatchTSTConfig, PatchTSTForPretraining >>> file = hf_hub_download( ... repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset" ... ) >>> batch = torch.load(file) >>> # Config for random mask pretraining >>> config = PatchTSTConfig( ... num_input_channels=7, ... context_length=512, ... patch_length=12, ... stride=12, ... mask_type='random', ... random_mask_ratio=0.4, ... use_cls_token=True, ... ) >>> # Config for forecast mask pretraining >>> config = PatchTSTConfig( ... num_input_channels=7, ... context_length=512, ... patch_length=12, ... stride=12, ... mask_type='forecast', ... num_forecast_mask_patches=5, ... use_cls_token=True, ... ) >>> model = PatchTSTForPretraining(config) >>> # during training, one provides both past and future values >>> outputs = model(past_values=batch["past_values"]) >>> loss = outputs.loss >>> loss.backward() ```NTrrrr:rnoneZ reductionr<r=rr)rrr5r)rrrr#rrMSELossrrrnrzr5rr) r(rrrr:r model_outputZx_hatrloss_valZ masked_lossrrr+r+r,rIs,E  $ zPatchTSTForPretraining.forward)NNNN)rJrKrLrrrCrQrrPrrrrIrRr+r+r)r,r"s&  r"cr)PatchTSTClassificationHeadrcsdt|j|_|j|_tjdd|_|jdkrt|jnt |_ t |j |j |j|_dSNrZ start_dimr)rrr pooling_typerFlattenflattenrrrrr#rrV num_targetsr rYr)r+r,rfs  z#PatchTSTClassificationHead.__init__r!cCs|jr|dddddddf}n"|jdkr|jdd}n|jdkr+|jddj}n td|jd||}|||}|S) a[ Parameters: embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*): Embedding from the model Returns: `torch.Tensor` of shape `(bs, num_targets)` Nrrr0r=rpooling operator  is not implemented yet) rr.rrvaluesr!r0r rr(r!pooled_embeddingr[r+r+r,rIns    z"PatchTSTClassificationHead.forwardrr+r+r)r,r+esr+z0 The PatchTST for classification model. csxeZdZdeffdd Ze     ddejdeejdee dee d ee d ee d e e e ffd d Z ZS)PatchTSTForClassificationrcsBt||jrtdd|_t||_t||_| dS)N+Setting `do_mask_input` parameter to False.F) rrrloggerwarningrrr+r#rrYr)r+r,rs     z"PatchTSTForClassification.__init__Nr target_valuesrrr:rr;c Cs|dur|n|jj}|j||||dd}||j}d} |dur)t} | ||} |sC|f|dd} | dur?| f| } | S| } | St| ||j|j dS)ac past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model target_values (`torch.Tensor`, *optional*): Labels associates with the `past_values` past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). Examples: ```python >>> from transformers import PatchTSTConfig, PatchTSTForClassification >>> # classification task with two input channel2 and 3 classes >>> config = PatchTSTConfig( ... num_input_channels=2, ... num_targets=3, ... context_length=512, ... patch_length=12, ... stride=12, ... use_cls_token=True, ... ) >>> model = PatchTSTForClassification(config=config) >>> # during inference, one only provides past values >>> past_values = torch.randn(20, 512, 2) >>> outputs = model(past_values=past_values) >>> labels = outputs.prediction_logits ```NTr$rr)rrr5r) rrrr#rrZCrossEntropyLossrr5r) r(rr;rrr:rr)y_hatr*rrr+r+r,rIs2,  z!PatchTSTForClassification.forwardr)rJrKrLrrrrCrQrrPrrrrIrRr+r+r)r,r7s. r7z, The PatchTST for regression Model. cs8eZdZd dedeffdd ZdejfddZZ S) PatchTSTPredictionHeadNrrcsDt|j|_|j|_|j|_|j|_|js|jr|j}n|j|}|jsvt|_ t|_ t|_ t |jD]8}|j tjdd|durW|j t||jn |j |||j |jdkrnt|jntq;dStjdd|_|durt||j|_n|||_|jdkrt|jnt|_dS)a num_patches (`int`): The number of patches in the input sequence. distribution_output (`DistributionOutput`, *optional*): The distribution output layer for probabilistic forecasting. If None, a linear output layer is used. r0r-Nr)rrshare_projectionrrr.rVrr projectionsdropoutsflattensrr|r/r#prediction_lengthget_parameter_projectionrrrr0 projectionr)r(rrdistribution_outputr rr)r+r,rs0      (  $zPatchTSTPredictionHead.__init__r!cCs|jr|dddddddf}n|jdkr|jdd}n|jdkr+|jddj}n|}|jseg}t|jD]%}|j||dd|ddf}|j ||}|j ||}| |q7t j |dd}n||}||}||}t|trtdd |D}|S|dd}|S) aj Parameters: embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*): Embedding from the model Returns: `torch.Tensor` of shape `(bs, forecast_len, num_channels)` Nrrr0r=rrcss|] }|ddVqdS)r0rN)r2)rszr+r+r,rErz1PatchTSTPredictionHead.forward..)rr.rrr4r>rrrAr@r?r|rCrr0rrDrxrr2)r(r!r6r[rr+r+r,rIs.          zPatchTSTPredictionHead.forwardr) rJrKrLrrNrrCrQrIrRr+r+r)r,r=s+r=z, The PatchTST for prediction model. cseZdZdeffdd Z     ddejdeejdeejdeed eed eed e e e ffd d Z ddejdeejd e fddZZS)PatchTSTForPredictionrcst||jrtdd|_t||_|jdkrd|_n/|jdkr,t |j d|_n"|jdkr9t |j d|_n|jdkrFt |j d|_nt d|jt||jjj|jd |_|dS) Nr8Fmse student_tr=normalnegative_binomialUnknown distribution output )rE)rrrr9r:rrrrEr rBr r r!r=rrr#rrYr)r+r,rQs$        zPatchTSTForPrediction.__init__Nrrrrr:rr;cCs|dur|n|jj}|j||||dd}||j}d} |jr"|} n||j|j} |durQ|jrF|jj||j|jd} t | |} t | } n t j dd} | | |} |j} |j}|sq| f|dd}| durm| f|}|S|}|St | | |j|j| |d S) aV Parameters: past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). future_values (`torch.Tensor` of shape `(bs, forecast_len, num_input_channels)`, *optional*): Future target values associated with the `past_values` output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers output_attentions (`bool`, *optional*): Whether or not to return the output attention of all layers return_dict (`bool`, *optional*): Whether or not to return a `ModelOutput` instead of a plain tuple. Returns: `PatchTSTForPredictionOutput` or tuple of `torch.Tensor` (if `return_dict`=False or `config.return_dict`=False) Examples: ```python >>> from huggingface_hub import hf_hub_download >>> import torch >>> from transformers import PatchTSTConfig, PatchTSTForPrediction >>> file = hf_hub_download( ... repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset" ... ) >>> batch = torch.load(file) >>> # Prediction task with 7 input channels and prediction length is 96 >>> model = PatchTSTForPrediction.from_pretrained("namctin/patchtst_etth1_forecast") >>> # during training, one provides both past and future values >>> outputs = model( ... past_values=batch["past_values"], ... future_values=batch["future_values"], ... ) >>> loss = outputs.loss >>> loss.backward() >>> # during inference, one only provides past values, the model outputs future values >>> outputs = model(past_values=batch["past_values"]) >>> prediction_outputs = outputs.prediction_outputs ```NTr$rrrr&rr<)rrr5rrr)rrrr#rrErr distributionrrrr(rr5r)r(rrrrr:rr)r<r*Z y_hat_outrNrrrrr+r+r,rInsL=      zPatchTSTForPrediction.forwardcsr|jj}||d|dd}|jr.|jj|j|j|jdfddt|D}tj |dd}n|j d}t |d S) a  Generate sequences of sample predictions from a model with a probability distribution head. Parameters: past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`): Past values of the time series that serves as context in order to predict the future. past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). Return: [`SamplePatchTSTOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of samples, prediction_length, 1)` or `(batch_size, number of samples, prediction_length, num_input_channels)` for multivariate predictions. NF)rrrrrMcg|]}qSr+samplerrrNr+r,rtrz2PatchTSTForPrediction.generate..rr=r) rnum_parallel_samplesrErNrrrrrCrrgrr(rrrTrZsamplesr+rRr,generates   zPatchTSTForPrediction.generaterr)rJrKrLrrrCrQrrPrrrrIrrVrRr+r+r)r,rGKs<   prGcs8eZdZdZd deffdd ZdejfddZZ S) PatchTSTRegressionHeadz Regression head Nrcst|j|_|j|_|j|_||_|j|j}t j dd|_ |j dkr,t |j nt |_|dur?t ||j|_dS|||_dSr,)rrZ output_rangey_rangerr.rErrVrr/r0rrrrr#r1rDrC)r(rrEr r)r+r,rs   zPatchTSTRegressionHead.__init__r!cCs|jr|dddddddf}n"|jdkr|jdd}n|jdkr+|jddj}n td|jd|||}||}|j du|j du@r_t ||j d |j d|j d}|S) aY Parameters: embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*): Embedding from the model Returns: `torch.Tensor` of shape `(bs, output_dim)` Nrrr0r=rr2r3r) rr.rrr4r!rr0rDrErXrCZsigmoidr5r+r+r,rI!s    (zPatchTSTRegressionHead.forwardrr\r+r+r)r,rW srWz, The PatchTST for regression model. cseZdZdeffdd Ze     ddejdeejdeejdee d ee d ee d e e e ffd d Z ddejdeejd efddZZS)PatchTSTForRegressionrcst||jrtdd|_t||_|jdkrd|_n/|jdkr,t |j d|_n"|jdkr9t |j d|_n|jdkrFt |j d|_nt d|jt||j|_|dS) Nr8FrHrIr=rJrKrL)rrrr9r:rrrrEr r1r r r!rWr#rrYr)r+r,rHs        zPatchTSTForRegression.__init__Nrr;rrr:rr;c s|dur|njj}j||||dd}|j}d} |durIjr>j|} tfdd|D}t| |} t | } n t j dd} | ||} |sc|f|dd } | dur_| f| } | S| } | St | ||j |jd S) a# past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model target_values (`torch.Tensor` of shape `(bs, num_input_channels)`): Target values associates with the `past_values` past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). Whether or not to return a `ModelOutput` instead of a plain tuple. Examples: ```python >>> from transformers import PatchTSTConfig, PatchTSTForRegression >>> # Regression task with 6 input channels and regress 2 targets >>> model = PatchTSTForRegression.from_pretrained("namctin/patchtst_etth1_regression") >>> # during inference, one only provides past values, the model outputs future values >>> past_values = torch.randn(20, 512, 6) >>> outputs = model(past_values=past_values) >>> regression_outputs = outputs.regression_outputs ```NTr$csg|] }|djjqS)r<)r1rr1)rsitemr(r+r,rtsz1PatchTSTForRegression.forward..rr&rr)rrr5r)rrrr#rrErNrrrrr(rr5r) r(rr;rrr:rr)r<rrNrr+r[r,rIbs<%      zPatchTSTForRegression.forwardcsb|jj}||d|dd}|j|jfddt|D}tj|ddd||jj }t |d S) a Generate sequences of sample predictions from a model with a probability distribution head. Parameters: past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`): Past values of the time series that serves as context in order to predict the future. past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). Return: [`SamplePatchTSTOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of samples, num_targets)`. NF)rr;rrcrOr+rPrrrRr+r,rtrz2PatchTSTForRegression.generate..rr=r<rS) rrTrErNrrrCrr1r1rrUr+rRr,rVs zPatchTSTForRegression.generaterr)rJrKrLrrrrCrQrrPrrrrIrrVrRr+r+r)r,rYBs>  LrY)rrrGr"rYr7)NFrrr)FrMr dataclassesrtypingrrrrCrZ activationsrZmodeling_outputsr Zmodeling_utilsr Ztime_series_utilsr r r utilsrrrZconfiguration_patchtstrZ get_loggerrJr9ModulerrTrQrOlistrPrNrprrrrrrrrrrrrrr distributions Distributionrrrrr r rrr"r+r7r=rGrWrY__all__r+r+r+r,s        =  D0< !$8>""$7po%W`<7