Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 Google LLC and HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch TimesFM model."""16 17import math18from collections.abc import Sequence19from dataclasses import dataclass20from typing import Callable, Optional, Union21 22import torch23import torch.nn as nn24import torch.nn.functional as F25 26from ...modeling_flash_attention_utils import FlashAttentionKwargs27from ...modeling_outputs import BaseModelOutput28from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel29from ...processing_utils import Unpack30from ...utils import auto_docstring, can_return_tuple, logging31from ..llama.modeling_llama import LlamaRMSNorm32from ..phi4_multimodal.modeling_phi4_multimodal import simple_eager_attention_forward33from .configuration_timesfm import TimesFmConfig34 35 36logger = logging.get_logger(__name__)37 38 39@dataclass40@auto_docstring41class TimesFmOutput(BaseModelOutput):42    r"""43    loc (`torch.Tensor` of shape `(batch_size, )`):44        The mean of the time series inputs.45    scale (`torch.Tensor` of shape `(batch_size,)`):46        The scale of the time series inputs.47    """48 49    loc: Optional[torch.Tensor] = None50    scale: Optional[torch.Tensor] = None51 52 53@dataclass54@auto_docstring55class TimesFmOutputForPrediction(BaseModelOutput):56    r"""57    mean_predictions (`torch.Tensor` of shape `(batch_size, sequence_length)`):58        The mean predictions of the time series.59    full_predictions (`torch.Tensor` of shape `(batch_size, sequence_length)`):60        The full predictions of the time series including the mean and the quantiles.61    loss (`torch.Tensor` of shape `(1,)`, *optional*, returned when `future_values` is provided):62        The loss of the TimesFM model.63    """64 65    mean_predictions: Optional[torch.Tensor] = None66    full_predictions: Optional[torch.Tensor] = None67    loss: Optional[Union[torch.Tensor, float]] = None68 69 70class TimesFmMLP(nn.Module):71    """Pax MLP in pytorch."""72 73    def __init__(self, config: TimesFmConfig):74        super().__init__()75        hidden_size = config.hidden_size76        intermediate_size = config.intermediate_size77 78        self.gate_proj = nn.Linear(hidden_size, intermediate_size)79        self.down_proj = nn.Linear(intermediate_size, hidden_size)80        self.layer_norm = nn.LayerNorm(normalized_shape=hidden_size, eps=1e-6)81 82    def forward(self, x, paddings=None):83        gate_inp = self.layer_norm(x)84        gate = self.gate_proj(gate_inp)85        gate = F.relu(gate)86        outputs = self.down_proj(gate)87        if paddings is not None:88            outputs = outputs * (1.0 - paddings[:, :, None])89        return outputs + x90 91 92class TimesFmResidualBlock(nn.Module):93    """TimesFM residual block."""94 95    def __init__(self, input_dims, hidden_dims, output_dims):96        super().__init__()97        self.input_dims = input_dims98        self.hidden_dims = hidden_dims99        self.output_dims = output_dims100 101        self.input_layer = nn.Linear(input_dims, hidden_dims)102        self.activation = nn.SiLU()103        self.output_layer = nn.Linear(hidden_dims, output_dims)104        self.residual_layer = nn.Linear(input_dims, output_dims)105 106    def forward(self, x):107        hidden = self.input_layer(x)108        hidden = self.activation(hidden)109        output = self.output_layer(hidden)110        residual = self.residual_layer(x)111        return output + residual112 113 114class TimesFmRMSNorm(LlamaRMSNorm):115    pass116 117 118class TimesFmPositionalEmbedding(nn.Module):119    """Generates position embedding for a given 1-d sequence."""120 121    def __init__(self, config: TimesFmConfig):122        super().__init__()123        min_timescale = config.min_timescale124        max_timescale = config.max_timescale125        self.embedding_dims = config.hidden_size126 127        num_timescales = self.embedding_dims // 2128        log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / max(num_timescales - 1, 1)129        self.register_buffer(130            "inv_timescales",131            min_timescale * torch.exp(torch.arange(num_timescales, dtype=torch.float32) * -log_timescale_increment),132        )133 134    def forward(self, seq_length=None, position=None):135        """Generates a Tensor of sinusoids with different frequencies.136 137        Args:138            seq_length: an optional Python int defining the output sequence length.139              if the `position` argument is specified.140            position: [B, seq_length], optional position for each token in the141              sequence, only required when the sequence is packed.142 143        Returns:144            [B, seqlen, D] if `position` is specified, else [1, seqlen, D]145        """146        if position is None and seq_length is None:147            raise ValueError("Either position or seq_length must be provided")148 149        if position is None:150            # [1, seqlen]151            position = torch.arange(seq_length, dtype=torch.float32, device=self.inv_timescales.device).unsqueeze(0)152        elif position.ndim != 2:153            raise ValueError(f"position must be 2-dimensional, got shape {position.shape}")154 155        scaled_time = position.view(*position.shape, 1) * self.inv_timescales.view(1, 1, -1)156        signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2)157 158        # Padding to ensure correct embedding dimension159        signal = F.pad(signal, (0, 0, 0, self.embedding_dims % 2))160        return signal161 162 163class TimesFmAttention(nn.Module):164    """Implements the attention used in TimesFM. One key difference is that there is _per_dim_scaling of the query."""165 166    def __init__(self, config: TimesFmConfig, layer_idx: int):167        super().__init__()168        self.config = config169        self.is_causal = True170        self.attention_dropout = config.attention_dropout171        self.layer_idx = layer_idx172 173        self.num_heads = config.num_attention_heads174        self.hidden_size = config.hidden_size175        self.head_dim = config.head_dim176 177        self.q_size = self.num_heads * self.head_dim178        self.kv_size = self.num_heads * self.head_dim179        self.scaling = nn.Parameter(torch.empty((self.head_dim,)))180 181        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)182        self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)183        self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)184        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size)185 186    def _scale_query(self, query: torch.Tensor) -> torch.Tensor:187        scale = F.softplus(self.scaling).mul(1.442695041 / math.sqrt(self.head_dim))188        return query * scale[None, None, None, :]189 190    def forward(191        self,192        hidden_states: torch.Tensor,193        attention_mask: Optional[torch.Tensor] = None,194        **kwargs: Unpack[FlashAttentionKwargs],195    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:196        input_shape = hidden_states.shape[:-1]197        hidden_shape = (*input_shape, -1, self.head_dim)198 199        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)200        query_states = self._scale_query(query_states)201        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)202        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)203 204        attention_interface: Callable = simple_eager_attention_forward205        if self.config._attn_implementation != "eager":206            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]207 208        attn_output, attn_weights = attention_interface(209            self,210            query_states,211            key_states,212            value_states,213            attention_mask,214            dropout=0.0 if not self.training else self.attention_dropout,215            scaling=1.0,216            **kwargs,217        )218        attn_output = attn_output.reshape(*input_shape, -1).contiguous()219        attn_output = self.o_proj(attn_output)220        return attn_output, attn_weights221 222 223class TimesFmDecoderLayer(nn.Module):224    """Transformer layer."""225 226    def __init__(self, config: TimesFmConfig, layer_idx: int):227        super().__init__()228 229        self.self_attn = TimesFmAttention(config, layer_idx=layer_idx)230        self.mlp = TimesFmMLP(config)231        self.input_layernorm = TimesFmRMSNorm(config.hidden_size, eps=config.rms_norm_eps)232 233    def forward(234        self,235        hidden_states: torch.Tensor,236        attention_mask: torch.Tensor,237        paddings: torch.Tensor,238        output_attentions: bool = False,239    ) -> tuple[Optional[torch.Tensor], torch.Tensor]:240        # Self Attention241        residual = hidden_states242        hidden_states = self.input_layernorm(hidden_states)243        hidden_states, scores = self.self_attn(244            hidden_states=hidden_states,245            attention_mask=attention_mask,246            output_attentions=output_attentions,247        )248        hidden_states = residual + hidden_states249 250        # MLP251        hidden_states = self.mlp(hidden_states, paddings=paddings)252 253        return scores, hidden_states254 255 256@auto_docstring257class TimesFmPreTrainedModel(PreTrainedModel):258    config: TimesFmConfig259    base_model_prefix = "timesfm"260    _no_split_modules = ["TimesFmDecoderLayer"]261    main_input_name = "past_values"262    _supports_sdpa = True263 264    def _init_weights(self, module):265        super()._init_weights(module)266        if isinstance(module, TimesFmAttention):267            # Initialize scaling parameter268            nn.init.ones_(module.scaling)269 270 271@auto_docstring272class TimesFmModel(TimesFmPreTrainedModel):273    def __init__(self, config: TimesFmConfig):274        super().__init__(config)275 276        self.config = config277        self.input_ff_layer = TimesFmResidualBlock(278            input_dims=2 * config.patch_length,279            output_dims=config.hidden_size,280            hidden_dims=config.intermediate_size,281        )282        self.freq_emb = nn.Embedding(num_embeddings=config.freq_size, embedding_dim=config.hidden_size)283        self.layers = nn.ModuleList(284            [TimesFmDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]285        )286        if self.config.use_positional_embedding:287            self.position_emb = TimesFmPositionalEmbedding(config=config)288 289        # Initialize weights and apply final processing290        self.post_init()291 292    def _forward_transform(293        self, inputs: torch.Tensor, patched_pads: torch.Tensor294    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:295        """Input is of shape [B, N, P]."""296        mu, sigma = self._timesfm_masked_mean_std(inputs, patched_pads)297        sigma = torch.where(298            sigma < self.config.tolerance,299            torch.tensor(1.0, dtype=sigma.dtype, device=sigma.device),300            sigma,301        )302 303        # Normalize each patch304        outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]305        outputs = torch.where(306            torch.abs(inputs - self.config.pad_val) < self.config.tolerance,307            torch.tensor(self.config.pad_val, dtype=outputs.dtype, device=outputs.device),308            outputs,309        )310        return outputs, (mu, sigma)311 312    @can_return_tuple313    @auto_docstring314    def forward(315        self,316        past_values: torch.Tensor,317        past_values_padding: torch.LongTensor,318        freq: torch.Tensor,319        output_attentions: bool = False,320        output_hidden_states: bool = False,321    ) -> TimesFmOutput:322        r"""323        past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):324            Past values of the time series that serves as input to the model.325        past_values_padding (`torch.LongTensor` of shape `(batch_size, sequence_length)`):326            The padding indicator of the time series.327        freq (`torch.LongTensor` of shape `(batch_size,)`):328            Frequency indices for the time series data.329        """330        # Reshape into patches (using view for efficiency)331        bsize = past_values.shape[0]332        patched_inputs = past_values.view(bsize, -1, self.config.patch_length)333        patched_pads = past_values_padding.view(bsize, -1, self.config.patch_length)334 335        patched_inputs = torch.where(336            torch.abs(patched_pads - 1.0) < self.config.tolerance,337            torch.tensor(0.0, dtype=patched_inputs.dtype, device=patched_inputs.device),338            patched_inputs,339        )340        patched_pads = torch.where(341            torch.abs(patched_inputs - self.config.pad_val) < self.config.tolerance,342            torch.tensor(1.0, dtype=patched_pads.dtype, device=patched_pads.device),343            patched_pads,344        )345        patched_inputs, stats = self._forward_transform(patched_inputs, patched_pads)346 347        # B x N x D348        patched_inputs = patched_inputs * (1.0 - patched_pads)349        concat_inputs = torch.cat([patched_inputs, patched_pads], dim=-1)350        model_input = self.input_ff_layer(concat_inputs)351 352        # A patch should not be padded even if there is at least one zero.353        patched_padding = torch.min(patched_pads, dim=-1)[0]  # Get the values from the min result354        if self.config.use_positional_embedding:355            pos_emb = self.position_emb(model_input.shape[1])356            pos_emb = torch.concat([pos_emb] * model_input.shape[0], dim=0)357            pos_emb = self._timesfm_shift_padded_seq(patched_padding, pos_emb)358            model_input += pos_emb359 360        f_emb = self.freq_emb(freq)  # B x 1 x D361        model_input += f_emb362 363        # Convert paddings to attention mask and combine with causal mask364        hidden_states = model_input365        attention_mask = self._prepare_4d_attention_mask(366            attention_mask=patched_padding,367            sequence_length=hidden_states.shape[1],368            dtype=hidden_states.dtype,369            device=hidden_states.device,370            is_causal=True,371        )372 373        all_attentions = []374        all_hidden_states = []375 376        for layer in self.layers[: self.config.num_hidden_layers]:377            scores, hidden_states = layer(378                hidden_states=hidden_states,379                attention_mask=attention_mask,380                paddings=patched_padding,381                output_attentions=output_attentions,382            )383            if output_attentions:384                all_attentions.append(scores)385            if output_hidden_states:386                all_hidden_states.append(hidden_states)387 388        if output_hidden_states:389            all_hidden_states = [model_input] + all_hidden_states390        else:391            all_hidden_states = None392 393        return TimesFmOutput(394            last_hidden_state=hidden_states,395            hidden_states=all_hidden_states,396            attentions=all_attentions if output_attentions else None,397            loc=stats[0],398            scale=stats[1],399        )400 401    @staticmethod402    def _prepare_4d_attention_mask(403        attention_mask: Optional[torch.Tensor],404        sequence_length: int,405        dtype: torch.dtype,406        device: torch.device,407        is_causal: bool = True,408    ) -> Optional[torch.Tensor]:409        """410        Creates 4D attention mask and combines causal and padding masks if needed.411 412        Args:413            attention_mask: Optional tensor of shape (batch_size, seq_length) containing padding mask414            sequence_length: Length of the sequence415            dtype: Data type of the mask416            device: Device of the mask417            is_causal: Whether to apply causal masking418 419        Returns:420            4D attention mask of shape (batch_size, 1, seq_length, seq_length)421        """422        # Get minimum value for the dtype423        min_value = torch.finfo(dtype).min if dtype.is_floating_point else torch.iinfo(dtype).min424 425        # Handle padding mask426        if attention_mask is not None:427            # Convert 2D padding mask to 4D attention mask428            attention_mask = attention_mask.view(attention_mask.shape[0], 1, 1, -1)429            attention_mask = attention_mask * min_value430 431        # Create causal mask if needed432        if is_causal:433            causal_mask = torch.triu(434                torch.ones((sequence_length, sequence_length), dtype=dtype, device=device) * min_value,435                diagonal=1,436            )437            causal_mask = causal_mask.view(1, 1, sequence_length, sequence_length)438 439            # Combine with padding mask if it exists440            if attention_mask is not None:441                attention_mask = torch.minimum(attention_mask, causal_mask)442            else:443                attention_mask = causal_mask444 445        return attention_mask446 447    @staticmethod448    def _timesfm_masked_mean_std(inputs: torch.Tensor, padding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:449        """Calculates mean and standard deviation of `inputs` across axis 1.450 451        It excludes values where `padding` is 1.452 453        Args:454            inputs: A PyTorch tensor of shape [b, n, p].455            padding: A PyTorch tensor of shape [b, n, p] with values 0 or 1.456 457        Returns:458            A tuple containing the mean and standard deviation.459            We return the statistics of the first patch with more than three non-padded values.460        """461 462        # Selecting the first patch with more than 3 unpadded values.463        def _get_patch_index(arr: torch.Tensor):464            indices = torch.argmax((arr >= 3).to(torch.int32), dim=1)465            row_sum = (arr >= 3).to(torch.int32).sum(dim=1)466            return torch.where(row_sum == 0, arr.shape[1] - 1, indices)467 468        pad_sum = torch.sum(1 - padding, dim=2)469        patch_indices = _get_patch_index(pad_sum)470        bidxs = torch.arange(inputs.shape[0])471 472        arr = inputs[bidxs, patch_indices, :]473        pad = padding[bidxs, patch_indices, :]474 475        # Create a mask where padding is 0476        mask = 1 - pad477 478        # Calculate the number of valid elements479        num_valid_elements = torch.sum(mask, dim=1)480        num_valid_elements = torch.where(481            num_valid_elements == 0,482            torch.tensor(1, dtype=num_valid_elements.dtype, device=num_valid_elements.device),483            num_valid_elements,484        )485 486        # Calculate the masked sum and squared sum487        masked_sum = torch.sum(arr * mask, dim=1)488        masked_squared_sum = torch.sum((arr * mask) ** 2, dim=1)489 490        # Calculate the masked mean and standard deviation491        masked_mean = masked_sum / num_valid_elements492        masked_var = masked_squared_sum / num_valid_elements - masked_mean**2493        masked_var = torch.where(494            masked_var < 0.0,495            torch.tensor(0.0, dtype=masked_var.dtype, device=masked_var.device),496            masked_var,497        )498        masked_std = torch.sqrt(masked_var)499 500        return masked_mean, masked_std501 502    @staticmethod503    def _timesfm_shift_padded_seq(mask: torch.Tensor, seq: torch.Tensor) -> torch.Tensor:504        """Shifts rows of seq based on the first 0 in each row of the mask.505 506        Args:507            mask: mask tensor of shape [B, N]508            seq: seq tensor of shape [B, N, P]509 510        Returns:511            The shifted sequence.512        """513        batch_size, num_seq, feature_dim = seq.shape514 515        new_mask: torch.BoolTensor = mask == 0516 517        # Use argmax to find the first True value in each row518        indices = new_mask.to(torch.int32).argmax(dim=1)519 520        # Handle rows with all zeros521        indices[~new_mask.any(dim=1)] = -1522 523        # Create index ranges for each sequence in the batch524        idx_range = torch.arange(num_seq, device=seq.device).view(1, -1, 1).expand(batch_size, -1, feature_dim)525 526        # Calculate shifted indices for each element in each sequence527        shifted_idx = (idx_range - indices[:, None, None]) % num_seq528 529        # Gather values from seq using shifted indices530        shifted_seq = seq.gather(1, shifted_idx)531 532        return shifted_seq533 534 535class TimesFmModelForPrediction(TimesFmPreTrainedModel):536    """TimesFM model for quantile and mean prediction."""537 538    def __init__(self, config: TimesFmConfig):539        super().__init__(config)540 541        self.config = config542        self.context_len = config.context_length543        self.horizon_len = config.horizon_length544 545        self.decoder = TimesFmModel(config)546 547        # quantile and mean output548        self.horizon_ff_layer = TimesFmResidualBlock(549            input_dims=config.hidden_size,550            output_dims=config.horizon_length * (1 + len(config.quantiles)),551            hidden_dims=config.intermediate_size,552        )553 554        # Initialize weights and apply final processing555        self.post_init()556 557    def _preprocess(558        self, inputs: Sequence[torch.Tensor], freq: Sequence[int]559    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:560        """Formats and pads raw inputs to feed into the model.561 562        This function both pads each time series to match the context length, and563        pads the inputs to meet the SPMD shape requirement.564 565        Args:566          inputs: A list of 1d Tensors. Each Tensor is the context time series of567            a single forecast task.568          freq: list of frequencies569 570        Returns:571        A tuple of:572        - the padded input time series to meet the model required context.573        - the padding indicator.574        - the number of padded examples for SPMD so that each core has the same575            number (a multiple of `batch_size`) of examples.576        """577        input_ts, input_padding, inp_freq = [], [], []578 579        for i, ts in enumerate(inputs):580            input_len = ts.shape[0]581            padding = torch.zeros(input_len + self.horizon_len, dtype=ts.dtype, device=ts.device)582            if input_len < self.context_len:583                num_front_pad = self.context_len - input_len584                ts = torch.cat([torch.zeros(num_front_pad, dtype=ts.dtype, device=ts.device), ts], dim=0)585                padding = torch.cat([torch.ones(num_front_pad, dtype=ts.dtype, device=padding.device), padding], dim=0)586            elif input_len > self.context_len:587                ts = ts[-self.context_len :]588                padding = padding[-(self.context_len + self.horizon_len) :]589 590            input_ts.append(ts)591            input_padding.append(padding)592            inp_freq.append(freq[i])593 594        return (595            torch.stack(input_ts, dim=0),596            torch.stack(input_padding, dim=0),597            torch.tensor(inp_freq, dtype=torch.int32).reshape(-1, 1),598        )599 600    def _postprocess_output(601        self, model_output: torch.Tensor, stats: tuple[torch.Tensor, torch.Tensor]602    ) -> torch.Tensor:603        """Postprocess output of stacked transformer."""604 605        # B x N x (H.Q)606        output_ts = self.horizon_ff_layer(model_output)607 608        # Reshape using view609        b, n, _ = output_ts.shape610        output_ts = output_ts.view(b, n, self.config.horizon_length, len(self.config.quantiles) + 1)611 612        mu, sigma = stats613        return output_ts * sigma[:, None, None, None] + mu[:, None, None, None]614 615    def _quantile_loss(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:616        losses = []617        for i, q in enumerate(self.config.quantiles):618            errors = targets - predictions[..., i]619            loss = torch.max((q - 1) * errors, q * errors)620            losses.append(loss.mean())621        return torch.stack(losses).mean()622 623    @can_return_tuple624    @auto_docstring625    def forward(626        self,627        past_values: Sequence[torch.Tensor],628        freq: Optional[Sequence[Union[torch.Tensor, int]]] = None,629        window_size: Optional[int] = None,630        future_values: Optional[torch.Tensor] = None,631        forecast_context_len: Optional[int] = None,632        return_forecast_on_context: bool = False,633        truncate_negative: bool = False,634        output_attentions: Optional[bool] = None,635        output_hidden_states: Optional[bool] = None,636    ) -> TimesFmOutputForPrediction:637        r"""638        past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):639            Past values of the time series that serves as input to the model.640        freq (`torch.LongTensor` of shape `(batch_size,)`):641            Frequency indices for the time series data.642        window_size (`int`, *optional*):643            Window size of trend + residual decomposition. If None then we do not do decomposition.644        future_values (`torch.Tensor`, *optional*):645            Optional future time series values to be used for loss computation.646        forecast_context_len (`int`, *optional*):647            Optional max context length.648        return_forecast_on_context (`bool`, *optional*):649            True to return the forecast on the context when available, i.e. after the first input patch.650        truncate_negative (`bool`, *optional*):651            Truncate to only non-negative values if any of the contexts have non-negative values,652            otherwise do nothing.653        output_attentions (`bool`, *optional*):654            Whether to output the attentions.655        output_hidden_states (`bool`, *optional*):656            Whether to output the hidden states.657 658        Example:659 660        ```python661        >>> from transformers import TimesFmModelForPrediction662 663        >>> model = TimesFmModelForPrediction.from_pretrained("google/timesfm-2.0-500m-pytorch")664 665        >>> forecast_input = [torch.linspace(0, 20, 100).sin(), torch.linspace(0, 20, 200).sin(), torch.linspace(0, 20, 400).sin()]666        >>> frequency_input = torch.tensor([0, 1, 2], dtype=torch.long)667 668        >>> # Generate669        >>> with torch.no_grad():670        >>>     outputs = model(past_values=forecast_input, freq=frequency_input, return_dict=True)671        >>>     point_forecast_conv = outputs.mean_predictions672        >>>     quantile_forecast_conv = outputs.full_predictions673        ```674        """675        if forecast_context_len is None:676            fcontext_len = self.context_len677        else:678            fcontext_len = forecast_context_len679 680        # Get device from first input tensor681        device = past_values[0].device682 683        # Truncate inputs to forecast_context_len684        inputs = [ts[-fcontext_len:] for ts in past_values]685        inp_min = torch.min(torch.stack([torch.min(ts) for ts in inputs]))686 687        if window_size is not None:688            new_inputs = []689            new_freqs = []690            for i, ts in enumerate(inputs):691                new_inputs.extend(self._timesfm_moving_average(ts, window_size))692                if freq is not None:693                    new_freqs.extend([freq[i]] * 2)694            inputs = new_inputs695            if freq is not None:696                freq = new_freqs697 698        if freq is None:699            logger.info("No frequency provided via `freq`. Default to high (0).")700            freq = [0] * len(inputs)701 702        if output_attentions is None:703            output_attentions = self.config.output_attentions704        if output_hidden_states is None:705            output_hidden_states = self.config.output_hidden_states706 707        input_ts, input_padding, inp_freq = self._preprocess(inputs, freq)708        # Move tensors to the same device as input709        input_ts = input_ts.to(device)710        input_padding = input_padding.to(device)711        inp_freq = inp_freq.to(device)712 713        final_out = input_ts714        context_len = final_out.shape[1]715        full_outputs = []716 717        if input_padding.shape[1] != final_out.shape[1] + self.horizon_len:718            raise ValueError(719                "Length of paddings must match length of input + horizon_len:"720                f" {input_padding.shape[1]} != {final_out.shape[1]} + {self.horizon_len}"721            )722        output_patch_len = self.config.horizon_length723 724        num_decode_patches = (self.horizon_len + output_patch_len - 1) // output_patch_len725        for step_index in range(num_decode_patches):726            current_padding = input_padding[:, 0 : final_out.shape[1]]727            input_ts = final_out[:, -fcontext_len:]728            input_padding = current_padding[:, -fcontext_len:]729            decoder_output = self.decoder(730                past_values=input_ts,731                past_values_padding=input_padding,732                freq=inp_freq,733                output_attentions=output_attentions,734                output_hidden_states=output_hidden_states,735            )736            fprop_outputs = self._postprocess_output(737                decoder_output.last_hidden_state,738                (decoder_output.loc, decoder_output.scale),739            )740 741            if return_forecast_on_context and step_index == 0:742                # For the first decodings step, collect the model forecast on the743                # context except the unavailable first input batch forecast.744                new_full_ts = fprop_outputs[:, :-1, : self.config.patch_length, :]745                # We have to use reshape and not view for non-contiguous memory746                new_full_ts = new_full_ts.reshape(new_full_ts.size(0), -1, new_full_ts.size(3))747 748                full_outputs.append(new_full_ts)749 750            # (full batch, last patch, output_patch_len, index of mean forecast = 0)751            new_ts = fprop_outputs[:, -1, :output_patch_len, 0]752            new_full_ts = fprop_outputs[:, -1, :output_patch_len, :]753            # (full batch, last patch, output_patch_len, all output indices)754            full_outputs.append(new_full_ts)755            final_out = torch.concatenate([final_out, new_ts], axis=-1)756 757        if return_forecast_on_context:758            # `full_outputs` indexing starts at after the first input patch.759            full_outputs = torch.concatenate(full_outputs, axis=1)[760                :, : (context_len - self.config.patch_length + self.horizon_len), :761            ]762        else:763            # `full_outputs` indexing starts at the forecast horizon.764            full_outputs = torch.concatenate(full_outputs, axis=1)[:, 0 : self.horizon_len, :]765 766        mean_outputs = full_outputs[:, :, 0]767        if window_size is not None:768            mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]769            full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]770        if inp_min >= 0 and truncate_negative:771            mean_outputs = torch.maximum(mean_outputs, 0.0)772            full_outputs = torch.maximum(full_outputs, 0.0)773 774        loss = None775        if future_values is not None:776            mse_loss = F.mse_loss(mean_outputs, future_values)777            quantile_loss = self._quantile_loss(full_outputs[:, :, 1:], future_values)778            loss = mse_loss + quantile_loss779 780        return TimesFmOutputForPrediction(781            last_hidden_state=decoder_output.last_hidden_state,782            attentions=decoder_output.attentions if output_attentions else None,783            hidden_states=decoder_output.hidden_states if output_hidden_states else None,784            mean_predictions=mean_outputs,785            full_predictions=full_outputs,786            loss=loss,787        )788 789    @staticmethod790    def _timesfm_moving_average(arr: torch.Tensor, window_size: int) -> list[torch.Tensor]:791        """Calculates the moving average using PyTorch's convolution function."""792        # Pad with zeros to handle initial window positions793        arr_padded = F.pad(arr, (window_size - 1, 0), "constant", 0)794        # Create a convolution kernel795        kernel = torch.ones(window_size, dtype=arr.dtype, device=arr.device) / window_size796        # Apply convolution to calculate the moving average797        smoothed_arr = F.conv1d(arr_padded.view(1, 1, -1), kernel.view(1, 1, -1)).squeeze()798        return [smoothed_arr, arr - smoothed_arr]799 800 801__all__ = ["TimesFmModelForPrediction", "TimesFmPreTrainedModel", "TimesFmModel"]802