Aluode/PerceptionLabPortable
0
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/timesfm/modular_timesfm.py.3# Do NOT edit this file manually as any edits will be overwritten by the generation of4# the file from the modular. If any change should be done, please apply the change to the5# modular_timesfm.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2025 Google LLC and HuggingFace Inc. team.9#10# Licensed under the Apache License, Version 2.0 (the "License");11# you may not use this file except in compliance with the License.12# You may obtain a copy of the License at13#14# http://www.apache.org/licenses/LICENSE-2.015#16# Unless required by applicable law or agreed to in writing, software17# distributed under the License is distributed on an "AS IS" BASIS,18# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.19# See the License for the specific language governing permissions and20# limitations under the License.21 22import math23from collections.abc import Sequence24from dataclasses import dataclass25from typing import Callable, Optional, Union26 27import torch28import torch.nn as nn29import torch.nn.functional as F30 31from ...integrations import use_kernel_forward_from_hub32from ...modeling_flash_attention_utils import FlashAttentionKwargs33from ...modeling_outputs import BaseModelOutput34from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel35from ...processing_utils import Unpack36from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging37from .configuration_timesfm import TimesFmConfig38 39 40logger = logging.get_logger(__name__)41 42 43@dataclass44@auto_docstring45class TimesFmOutput(BaseModelOutput):46 r"""47 loc (`torch.Tensor` of shape `(batch_size, )`):48 The mean of the time series inputs.49 scale (`torch.Tensor` of shape `(batch_size,)`):50 The scale of the time series inputs.51 """52 53 loc: Optional[torch.Tensor] = None54 scale: Optional[torch.Tensor] = None55 56 57@dataclass58@auto_docstring59class TimesFmOutputForPrediction(BaseModelOutput):60 r"""61 mean_predictions (`torch.Tensor` of shape `(batch_size, sequence_length)`):62 The mean predictions of the time series.63 full_predictions (`torch.Tensor` of shape `(batch_size, sequence_length)`):64 The full predictions of the time series including the mean and the quantiles.65 loss (`torch.Tensor` of shape `(1,)`, *optional*, returned when `future_values` is provided):66 The loss of the TimesFM model.67 """68 69 mean_predictions: Optional[torch.Tensor] = None70 full_predictions: Optional[torch.Tensor] = None71 loss: Optional[Union[torch.Tensor, float]] = None72 73 74class TimesFmMLP(nn.Module):75 """Pax MLP in pytorch."""76 77 def __init__(self, config: TimesFmConfig):78 super().__init__()79 hidden_size = config.hidden_size80 intermediate_size = config.intermediate_size81 82 self.gate_proj = nn.Linear(hidden_size, intermediate_size)83 self.down_proj = nn.Linear(intermediate_size, hidden_size)84 self.layer_norm = nn.LayerNorm(normalized_shape=hidden_size, eps=1e-6)85 86 def forward(self, x, paddings=None):87 gate_inp = self.layer_norm(x)88 gate = self.gate_proj(gate_inp)89 gate = F.relu(gate)90 outputs = self.down_proj(gate)91 if paddings is not None:92 outputs = outputs * (1.0 - paddings[:, :, None])93 return outputs + x94 95 96class TimesFmResidualBlock(nn.Module):97 """TimesFM residual block."""98 99 def __init__(self, input_dims, hidden_dims, output_dims):100 super().__init__()101 self.input_dims = input_dims102 self.hidden_dims = hidden_dims103 self.output_dims = output_dims104 105 self.input_layer = nn.Linear(input_dims, hidden_dims)106 self.activation = nn.SiLU()107 self.output_layer = nn.Linear(hidden_dims, output_dims)108 self.residual_layer = nn.Linear(input_dims, output_dims)109 110 def forward(self, x):111 hidden = self.input_layer(x)112 hidden = self.activation(hidden)113 output = self.output_layer(hidden)114 residual = self.residual_layer(x)115 return output + residual116 117 118@use_kernel_forward_from_hub("RMSNorm")119class TimesFmRMSNorm(nn.Module):120 def __init__(self, hidden_size, eps=1e-6):121 """122 TimesFmRMSNorm is equivalent to T5LayerNorm123 """124 super().__init__()125 self.weight = nn.Parameter(torch.ones(hidden_size))126 self.variance_epsilon = eps127 128 def forward(self, hidden_states):129 input_dtype = hidden_states.dtype130 hidden_states = hidden_states.to(torch.float32)131 variance = hidden_states.pow(2).mean(-1, keepdim=True)132 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)133 return self.weight * hidden_states.to(input_dtype)134 135 def extra_repr(self):136 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"137 138 139class TimesFmPositionalEmbedding(nn.Module):140 """Generates position embedding for a given 1-d sequence."""141 142 def __init__(self, config: TimesFmConfig):143 super().__init__()144 min_timescale = config.min_timescale145 max_timescale = config.max_timescale146 self.embedding_dims = config.hidden_size147 148 num_timescales = self.embedding_dims // 2149 log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / max(num_timescales - 1, 1)150 self.register_buffer(151 "inv_timescales",152 min_timescale * torch.exp(torch.arange(num_timescales, dtype=torch.float32) * -log_timescale_increment),153 )154 155 def forward(self, seq_length=None, position=None):156 """Generates a Tensor of sinusoids with different frequencies.157 158 Args:159 seq_length: an optional Python int defining the output sequence length.160 if the `position` argument is specified.161 position: [B, seq_length], optional position for each token in the162 sequence, only required when the sequence is packed.163 164 Returns:165 [B, seqlen, D] if `position` is specified, else [1, seqlen, D]166 """167 if position is None and seq_length is None:168 raise ValueError("Either position or seq_length must be provided")169 170 if position is None:171 # [1, seqlen]172 position = torch.arange(seq_length, dtype=torch.float32, device=self.inv_timescales.device).unsqueeze(0)173 elif position.ndim != 2:174 raise ValueError(f"position must be 2-dimensional, got shape {position.shape}")175 176 scaled_time = position.view(*position.shape, 1) * self.inv_timescales.view(1, 1, -1)177 signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2)178 179 # Padding to ensure correct embedding dimension180 signal = F.pad(signal, (0, 0, 0, self.embedding_dims % 2))181 return signal182 183 184def simple_eager_attention_forward(185 module: nn.Module,186 query_states: torch.Tensor,187 key_states: torch.Tensor,188 value_states: torch.Tensor,189 attention_mask: Optional[torch.Tensor],190 scaling: float,191 dropout: float = 0.0,192 **kwargs: Unpack[TransformersKwargs],193):194 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * scaling195 if attention_mask is not None:196 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]197 attn_weights = attn_weights + causal_mask198 199 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)200 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)201 attn_output = torch.matmul(attn_weights, value_states)202 attn_output = attn_output.transpose(1, 2).contiguous()203 204 return attn_output, attn_weights205 206 207class TimesFmAttention(nn.Module):208 """Implements the attention used in TimesFM. One key difference is that there is _per_dim_scaling of the query."""209 210 def __init__(self, config: TimesFmConfig, layer_idx: int):211 super().__init__()212 self.config = config213 self.is_causal = True214 self.attention_dropout = config.attention_dropout215 self.layer_idx = layer_idx216 217 self.num_heads = config.num_attention_heads218 self.hidden_size = config.hidden_size219 self.head_dim = config.head_dim220 221 self.q_size = self.num_heads * self.head_dim222 self.kv_size = self.num_heads * self.head_dim223 self.scaling = nn.Parameter(torch.empty((self.head_dim,)))224 225 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)226 self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)227 self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim)228 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size)229 230 def _scale_query(self, query: torch.Tensor) -> torch.Tensor:231 scale = F.softplus(self.scaling).mul(1.442695041 / math.sqrt(self.head_dim))232 return query * scale[None, None, None, :]233 234 def forward(235 self,236 hidden_states: torch.Tensor,237 attention_mask: Optional[torch.Tensor] = None,238 **kwargs: Unpack[FlashAttentionKwargs],239 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:240 input_shape = hidden_states.shape[:-1]241 hidden_shape = (*input_shape, -1, self.head_dim)242 243 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)244 query_states = self._scale_query(query_states)245 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)246 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)247 248 attention_interface: Callable = simple_eager_attention_forward249 if self.config._attn_implementation != "eager":250 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]251 252 attn_output, attn_weights = attention_interface(253 self,254 query_states,255 key_states,256 value_states,257 attention_mask,258 dropout=0.0 if not self.training else self.attention_dropout,259 scaling=1.0,260 **kwargs,261 )262 attn_output = attn_output.reshape(*input_shape, -1).contiguous()263 attn_output = self.o_proj(attn_output)264 return attn_output, attn_weights265 266 267class TimesFmDecoderLayer(nn.Module):268 """Transformer layer."""269 270 def __init__(self, config: TimesFmConfig, layer_idx: int):271 super().__init__()272 273 self.self_attn = TimesFmAttention(config, layer_idx=layer_idx)274 self.mlp = TimesFmMLP(config)275 self.input_layernorm = TimesFmRMSNorm(config.hidden_size, eps=config.rms_norm_eps)276 277 def forward(278 self,279 hidden_states: torch.Tensor,280 attention_mask: torch.Tensor,281 paddings: torch.Tensor,282 output_attentions: bool = False,283 ) -> tuple[Optional[torch.Tensor], torch.Tensor]:284 # Self Attention285 residual = hidden_states286 hidden_states = self.input_layernorm(hidden_states)287 hidden_states, scores = self.self_attn(288 hidden_states=hidden_states,289 attention_mask=attention_mask,290 output_attentions=output_attentions,291 )292 hidden_states = residual + hidden_states293 294 # MLP295 hidden_states = self.mlp(hidden_states, paddings=paddings)296 297 return scores, hidden_states298 299 300@auto_docstring301class TimesFmPreTrainedModel(PreTrainedModel):302 config: TimesFmConfig303 base_model_prefix = "timesfm"304 _no_split_modules = ["TimesFmDecoderLayer"]305 main_input_name = "past_values"306 _supports_sdpa = True307 308 def _init_weights(self, module):309 super()._init_weights(module)310 if isinstance(module, TimesFmAttention):311 # Initialize scaling parameter312 nn.init.ones_(module.scaling)313 314 315@auto_docstring316class TimesFmModel(TimesFmPreTrainedModel):317 def __init__(self, config: TimesFmConfig):318 super().__init__(config)319 320 self.config = config321 self.input_ff_layer = TimesFmResidualBlock(322 input_dims=2 * config.patch_length,323 output_dims=config.hidden_size,324 hidden_dims=config.intermediate_size,325 )326 self.freq_emb = nn.Embedding(num_embeddings=config.freq_size, embedding_dim=config.hidden_size)327 self.layers = nn.ModuleList(328 [TimesFmDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]329 )330 if self.config.use_positional_embedding:331 self.position_emb = TimesFmPositionalEmbedding(config=config)332 333 # Initialize weights and apply final processing334 self.post_init()335 336 def _forward_transform(337 self, inputs: torch.Tensor, patched_pads: torch.Tensor338 ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:339 """Input is of shape [B, N, P]."""340 mu, sigma = self._timesfm_masked_mean_std(inputs, patched_pads)341 sigma = torch.where(342 sigma < self.config.tolerance,343 torch.tensor(1.0, dtype=sigma.dtype, device=sigma.device),344 sigma,345 )346 347 # Normalize each patch348 outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]349 outputs = torch.where(350 torch.abs(inputs - self.config.pad_val) < self.config.tolerance,351 torch.tensor(self.config.pad_val, dtype=outputs.dtype, device=outputs.device),352 outputs,353 )354 return outputs, (mu, sigma)355 356 @can_return_tuple357 @auto_docstring358 def forward(359 self,360 past_values: torch.Tensor,361 past_values_padding: torch.LongTensor,362 freq: torch.Tensor,363 output_attentions: bool = False,364 output_hidden_states: bool = False,365 ) -> TimesFmOutput:366 r"""367 past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):368 Past values of the time series that serves as input to the model.369 past_values_padding (`torch.LongTensor` of shape `(batch_size, sequence_length)`):370 The padding indicator of the time series.371 freq (`torch.LongTensor` of shape `(batch_size,)`):372 Frequency indices for the time series data.373 """374 # Reshape into patches (using view for efficiency)375 bsize = past_values.shape[0]376 patched_inputs = past_values.view(bsize, -1, self.config.patch_length)377 patched_pads = past_values_padding.view(bsize, -1, self.config.patch_length)378 379 patched_inputs = torch.where(380 torch.abs(patched_pads - 1.0) < self.config.tolerance,381 torch.tensor(0.0, dtype=patched_inputs.dtype, device=patched_inputs.device),382 patched_inputs,383 )384 patched_pads = torch.where(385 torch.abs(patched_inputs - self.config.pad_val) < self.config.tolerance,386 torch.tensor(1.0, dtype=patched_pads.dtype, device=patched_pads.device),387 patched_pads,388 )389 patched_inputs, stats = self._forward_transform(patched_inputs, patched_pads)390 391 # B x N x D392 patched_inputs = patched_inputs * (1.0 - patched_pads)393 concat_inputs = torch.cat([patched_inputs, patched_pads], dim=-1)394 model_input = self.input_ff_layer(concat_inputs)395 396 # A patch should not be padded even if there is at least one zero.397 patched_padding = torch.min(patched_pads, dim=-1)[0] # Get the values from the min result398 if self.config.use_positional_embedding:399 pos_emb = self.position_emb(model_input.shape[1])400 pos_emb = torch.concat([pos_emb] * model_input.shape[0], dim=0)401 pos_emb = self._timesfm_shift_padded_seq(patched_padding, pos_emb)402 model_input += pos_emb403 404 f_emb = self.freq_emb(freq) # B x 1 x D405 model_input += f_emb406 407 # Convert paddings to attention mask and combine with causal mask408 hidden_states = model_input409 attention_mask = self._prepare_4d_attention_mask(410 attention_mask=patched_padding,411 sequence_length=hidden_states.shape[1],412 dtype=hidden_states.dtype,413 device=hidden_states.device,414 is_causal=True,415 )416 417 all_attentions = []418 all_hidden_states = []419 420 for layer in self.layers[: self.config.num_hidden_layers]:421 scores, hidden_states = layer(422 hidden_states=hidden_states,423 attention_mask=attention_mask,424 paddings=patched_padding,425 output_attentions=output_attentions,426 )427 if output_attentions:428 all_attentions.append(scores)429 if output_hidden_states:430 all_hidden_states.append(hidden_states)431 432 if output_hidden_states:433 all_hidden_states = [model_input] + all_hidden_states434 else:435 all_hidden_states = None436 437 return TimesFmOutput(438 last_hidden_state=hidden_states,439 hidden_states=all_hidden_states,440 attentions=all_attentions if output_attentions else None,441 loc=stats[0],442 scale=stats[1],443 )444 445 @staticmethod446 def _prepare_4d_attention_mask(447 attention_mask: Optional[torch.Tensor],448 sequence_length: int,449 dtype: torch.dtype,450 device: torch.device,451 is_causal: bool = True,452 ) -> Optional[torch.Tensor]:453 """454 Creates 4D attention mask and combines causal and padding masks if needed.455 456 Args:457 attention_mask: Optional tensor of shape (batch_size, seq_length) containing padding mask458 sequence_length: Length of the sequence459 dtype: Data type of the mask460 device: Device of the mask461 is_causal: Whether to apply causal masking462 463 Returns:464 4D attention mask of shape (batch_size, 1, seq_length, seq_length)465 """466 # Get minimum value for the dtype467 min_value = torch.finfo(dtype).min if dtype.is_floating_point else torch.iinfo(dtype).min468 469 # Handle padding mask470 if attention_mask is not None:471 # Convert 2D padding mask to 4D attention mask472 attention_mask = attention_mask.view(attention_mask.shape[0], 1, 1, -1)473 attention_mask = attention_mask * min_value474 475 # Create causal mask if needed476 if is_causal:477 causal_mask = torch.triu(478 torch.ones((sequence_length, sequence_length), dtype=dtype, device=device) * min_value,479 diagonal=1,480 )481 causal_mask = causal_mask.view(1, 1, sequence_length, sequence_length)482 483 # Combine with padding mask if it exists484 if attention_mask is not None:485 attention_mask = torch.minimum(attention_mask, causal_mask)486 else:487 attention_mask = causal_mask488 489 return attention_mask490 491 @staticmethod492 def _timesfm_masked_mean_std(inputs: torch.Tensor, padding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:493 """Calculates mean and standard deviation of `inputs` across axis 1.494 495 It excludes values where `padding` is 1.496 497 Args:498 inputs: A PyTorch tensor of shape [b, n, p].499 padding: A PyTorch tensor of shape [b, n, p] with values 0 or 1.500 501 Returns:502 A tuple containing the mean and standard deviation.503 We return the statistics of the first patch with more than three non-padded values.504 """505 506 # Selecting the first patch with more than 3 unpadded values.507 def _get_patch_index(arr: torch.Tensor):508 indices = torch.argmax((arr >= 3).to(torch.int32), dim=1)509 row_sum = (arr >= 3).to(torch.int32).sum(dim=1)510 return torch.where(row_sum == 0, arr.shape[1] - 1, indices)511 512 pad_sum = torch.sum(1 - padding, dim=2)513 patch_indices = _get_patch_index(pad_sum)514 bidxs = torch.arange(inputs.shape[0])515 516 arr = inputs[bidxs, patch_indices, :]517 pad = padding[bidxs, patch_indices, :]518 519 # Create a mask where padding is 0520 mask = 1 - pad521 522 # Calculate the number of valid elements523 num_valid_elements = torch.sum(mask, dim=1)524 num_valid_elements = torch.where(525 num_valid_elements == 0,526 torch.tensor(1, dtype=num_valid_elements.dtype, device=num_valid_elements.device),527 num_valid_elements,528 )529 530 # Calculate the masked sum and squared sum531 masked_sum = torch.sum(arr * mask, dim=1)532 masked_squared_sum = torch.sum((arr * mask) ** 2, dim=1)533 534 # Calculate the masked mean and standard deviation535 masked_mean = masked_sum / num_valid_elements536 masked_var = masked_squared_sum / num_valid_elements - masked_mean**2537 masked_var = torch.where(538 masked_var < 0.0,539 torch.tensor(0.0, dtype=masked_var.dtype, device=masked_var.device),540 masked_var,541 )542 masked_std = torch.sqrt(masked_var)543 544 return masked_mean, masked_std545 546 @staticmethod547 def _timesfm_shift_padded_seq(mask: torch.Tensor, seq: torch.Tensor) -> torch.Tensor:548 """Shifts rows of seq based on the first 0 in each row of the mask.549 550 Args:551 mask: mask tensor of shape [B, N]552 seq: seq tensor of shape [B, N, P]553 554 Returns:555 The shifted sequence.556 """557 batch_size, num_seq, feature_dim = seq.shape558 559 new_mask: torch.BoolTensor = mask == 0560 561 # Use argmax to find the first True value in each row562 indices = new_mask.to(torch.int32).argmax(dim=1)563 564 # Handle rows with all zeros565 indices[~new_mask.any(dim=1)] = -1566 567 # Create index ranges for each sequence in the batch568 idx_range = torch.arange(num_seq, device=seq.device).view(1, -1, 1).expand(batch_size, -1, feature_dim)569 570 # Calculate shifted indices for each element in each sequence571 shifted_idx = (idx_range - indices[:, None, None]) % num_seq572 573 # Gather values from seq using shifted indices574 shifted_seq = seq.gather(1, shifted_idx)575 576 return shifted_seq577 578 579class TimesFmModelForPrediction(TimesFmPreTrainedModel):580 """TimesFM model for quantile and mean prediction."""581 582 def __init__(self, config: TimesFmConfig):583 super().__init__(config)584 585 self.config = config586 self.context_len = config.context_length587 self.horizon_len = config.horizon_length588 589 self.decoder = TimesFmModel(config)590 591 # quantile and mean output592 self.horizon_ff_layer = TimesFmResidualBlock(593 input_dims=config.hidden_size,594 output_dims=config.horizon_length * (1 + len(config.quantiles)),595 hidden_dims=config.intermediate_size,596 )597 598 # Initialize weights and apply final processing599 self.post_init()600 601 def _preprocess(602 self, inputs: Sequence[torch.Tensor], freq: Sequence[int]603 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:604 """Formats and pads raw inputs to feed into the model.605 606 This function both pads each time series to match the context length, and607 pads the inputs to meet the SPMD shape requirement.608 609 Args:610 inputs: A list of 1d Tensors. Each Tensor is the context time series of611 a single forecast task.612 freq: list of frequencies613 614 Returns:615 A tuple of:616 - the padded input time series to meet the model required context.617 - the padding indicator.618 - the number of padded examples for SPMD so that each core has the same619 number (a multiple of `batch_size`) of examples.620 """621 input_ts, input_padding, inp_freq = [], [], []622 623 for i, ts in enumerate(inputs):624 input_len = ts.shape[0]625 padding = torch.zeros(input_len + self.horizon_len, dtype=ts.dtype, device=ts.device)626 if input_len < self.context_len:627 num_front_pad = self.context_len - input_len628 ts = torch.cat([torch.zeros(num_front_pad, dtype=ts.dtype, device=ts.device), ts], dim=0)629 padding = torch.cat([torch.ones(num_front_pad, dtype=ts.dtype, device=padding.device), padding], dim=0)630 elif input_len > self.context_len:631 ts = ts[-self.context_len :]632 padding = padding[-(self.context_len + self.horizon_len) :]633 634 input_ts.append(ts)635 input_padding.append(padding)636 inp_freq.append(freq[i])637 638 return (639 torch.stack(input_ts, dim=0),640 torch.stack(input_padding, dim=0),641 torch.tensor(inp_freq, dtype=torch.int32).reshape(-1, 1),642 )643 644 def _postprocess_output(645 self, model_output: torch.Tensor, stats: tuple[torch.Tensor, torch.Tensor]646 ) -> torch.Tensor:647 """Postprocess output of stacked transformer."""648 649 # B x N x (H.Q)650 output_ts = self.horizon_ff_layer(model_output)651 652 # Reshape using view653 b, n, _ = output_ts.shape654 output_ts = output_ts.view(b, n, self.config.horizon_length, len(self.config.quantiles) + 1)655 656 mu, sigma = stats657 return output_ts * sigma[:, None, None, None] + mu[:, None, None, None]658 659 def _quantile_loss(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:660 losses = []661 for i, q in enumerate(self.config.quantiles):662 errors = targets - predictions[..., i]663 loss = torch.max((q - 1) * errors, q * errors)664 losses.append(loss.mean())665 return torch.stack(losses).mean()666 667 @can_return_tuple668 @auto_docstring669 def forward(670 self,671 past_values: Sequence[torch.Tensor],672 freq: Optional[Sequence[Union[torch.Tensor, int]]] = None,673 window_size: Optional[int] = None,674 future_values: Optional[torch.Tensor] = None,675 forecast_context_len: Optional[int] = None,676 return_forecast_on_context: bool = False,677 truncate_negative: bool = False,678 output_attentions: Optional[bool] = None,679 output_hidden_states: Optional[bool] = None,680 ) -> TimesFmOutputForPrediction:681 r"""682 past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):683 Past values of the time series that serves as input to the model.684 freq (`torch.LongTensor` of shape `(batch_size,)`):685 Frequency indices for the time series data.686 window_size (`int`, *optional*):687 Window size of trend + residual decomposition. If None then we do not do decomposition.688 future_values (`torch.Tensor`, *optional*):689 Optional future time series values to be used for loss computation.690 forecast_context_len (`int`, *optional*):691 Optional max context length.692 return_forecast_on_context (`bool`, *optional*):693 True to return the forecast on the context when available, i.e. after the first input patch.694 truncate_negative (`bool`, *optional*):695 Truncate to only non-negative values if any of the contexts have non-negative values,696 otherwise do nothing.697 output_attentions (`bool`, *optional*):698 Whether to output the attentions.699 output_hidden_states (`bool`, *optional*):700 Whether to output the hidden states.701 702 Example:703 704 ```python705 >>> from transformers import TimesFmModelForPrediction706 707 >>> model = TimesFmModelForPrediction.from_pretrained("google/timesfm-2.0-500m-pytorch")708 709 >>> forecast_input = [torch.linspace(0, 20, 100).sin(), torch.linspace(0, 20, 200).sin(), torch.linspace(0, 20, 400).sin()]710 >>> frequency_input = torch.tensor([0, 1, 2], dtype=torch.long)711 712 >>> # Generate713 >>> with torch.no_grad():714 >>> outputs = model(past_values=forecast_input, freq=frequency_input, return_dict=True)715 >>> point_forecast_conv = outputs.mean_predictions716 >>> quantile_forecast_conv = outputs.full_predictions717 ```718 """719 if forecast_context_len is None:720 fcontext_len = self.context_len721 else:722 fcontext_len = forecast_context_len723 724 # Get device from first input tensor725 device = past_values[0].device726 727 # Truncate inputs to forecast_context_len728 inputs = [ts[-fcontext_len:] for ts in past_values]729 inp_min = torch.min(torch.stack([torch.min(ts) for ts in inputs]))730 731 if window_size is not None:732 new_inputs = []733 new_freqs = []734 for i, ts in enumerate(inputs):735 new_inputs.extend(self._timesfm_moving_average(ts, window_size))736 if freq is not None:737 new_freqs.extend([freq[i]] * 2)738 inputs = new_inputs739 if freq is not None:740 freq = new_freqs741 742 if freq is None:743 logger.info("No frequency provided via `freq`. Default to high (0).")744 freq = [0] * len(inputs)745 746 if output_attentions is None:747 output_attentions = self.config.output_attentions748 if output_hidden_states is None:749 output_hidden_states = self.config.output_hidden_states750 751 input_ts, input_padding, inp_freq = self._preprocess(inputs, freq)752 # Move tensors to the same device as input753 input_ts = input_ts.to(device)754 input_padding = input_padding.to(device)755 inp_freq = inp_freq.to(device)756 757 final_out = input_ts758 context_len = final_out.shape[1]759 full_outputs = []760 761 if input_padding.shape[1] != final_out.shape[1] + self.horizon_len:762 raise ValueError(763 "Length of paddings must match length of input + horizon_len:"764 f" {input_padding.shape[1]} != {final_out.shape[1]} + {self.horizon_len}"765 )766 output_patch_len = self.config.horizon_length767 768 num_decode_patches = (self.horizon_len + output_patch_len - 1) // output_patch_len769 for step_index in range(num_decode_patches):770 current_padding = input_padding[:, 0 : final_out.shape[1]]771 input_ts = final_out[:, -fcontext_len:]772 input_padding = current_padding[:, -fcontext_len:]773 decoder_output = self.decoder(774 past_values=input_ts,775 past_values_padding=input_padding,776 freq=inp_freq,777 output_attentions=output_attentions,778 output_hidden_states=output_hidden_states,779 )780 fprop_outputs = self._postprocess_output(781 decoder_output.last_hidden_state,782 (decoder_output.loc, decoder_output.scale),783 )784 785 if return_forecast_on_context and step_index == 0:786 # For the first decodings step, collect the model forecast on the787 # context except the unavailable first input batch forecast.788 new_full_ts = fprop_outputs[:, :-1, : self.config.patch_length, :]789 # We have to use reshape and not view for non-contiguous memory790 new_full_ts = new_full_ts.reshape(new_full_ts.size(0), -1, new_full_ts.size(3))791 792 full_outputs.append(new_full_ts)793 794 # (full batch, last patch, output_patch_len, index of mean forecast = 0)795 new_ts = fprop_outputs[:, -1, :output_patch_len, 0]796 new_full_ts = fprop_outputs[:, -1, :output_patch_len, :]797 # (full batch, last patch, output_patch_len, all output indices)798 full_outputs.append(new_full_ts)799 final_out = torch.concatenate([final_out, new_ts], axis=-1)800 801 if return_forecast_on_context:802 # `full_outputs` indexing starts at after the first input patch.803 full_outputs = torch.concatenate(full_outputs, axis=1)[804 :, : (context_len - self.config.patch_length + self.horizon_len), :805 ]806 else:807 # `full_outputs` indexing starts at the forecast horizon.808 full_outputs = torch.concatenate(full_outputs, axis=1)[:, 0 : self.horizon_len, :]809 810 mean_outputs = full_outputs[:, :, 0]811 if window_size is not None:812 mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]813 full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]814 if inp_min >= 0 and truncate_negative:815 mean_outputs = torch.maximum(mean_outputs, 0.0)816 full_outputs = torch.maximum(full_outputs, 0.0)817 818 loss = None819 if future_values is not None:820 mse_loss = F.mse_loss(mean_outputs, future_values)821 quantile_loss = self._quantile_loss(full_outputs[:, :, 1:], future_values)822 loss = mse_loss + quantile_loss823 824 return TimesFmOutputForPrediction(825 last_hidden_state=decoder_output.last_hidden_state,826 attentions=decoder_output.attentions if output_attentions else None,827 hidden_states=decoder_output.hidden_states if output_hidden_states else None,828 mean_predictions=mean_outputs,829 full_predictions=full_outputs,830 loss=loss,831 )832 833 @staticmethod834 def _timesfm_moving_average(arr: torch.Tensor, window_size: int) -> list[torch.Tensor]:835 """Calculates the moving average using PyTorch's convolution function."""836 # Pad with zeros to handle initial window positions837 arr_padded = F.pad(arr, (window_size - 1, 0), "constant", 0)838 # Create a convolution kernel839 kernel = torch.ones(window_size, dtype=arr.dtype, device=arr.device) / window_size840 # Apply convolution to calculate the moving average841 smoothed_arr = F.conv1d(arr_padded.view(1, 1, -1), kernel.view(1, 1, -1)).squeeze()842 return [smoothed_arr, arr - smoothed_arr]843 844 845__all__ = ["TimesFmModelForPrediction", "TimesFmPreTrainedModel", "TimesFmModel"]846