From c42a74091710ad96d0832de4263a962bed0fc83f Mon Sep 17 00:00:00 2001 From: Rajat Sen Date: Fri, 26 Sep 2025 22:47:56 +0000 Subject: [PATCH 1/3] loading the model using from_pretrained method --- pyproject.toml | 4 + src/timesfm/timesfm_2p5/timesfm_2p5_torch.py | 284 +++++++++---------- 2 files changed, 136 insertions(+), 152 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6055583..b11c5d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,10 @@ huggingface_hub = { version = ">=0.23.0", extras = ["cli"] } safetensors = ">=0.5.3" torch = { version = ">=2.0.0", extras = ["cuda"] } +[tool.ruff] +line-length = 100 +indent-width = 2 + [build-system] requires = ["poetry-core"] diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index dfb21aa..e82a0f7 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py @@ -11,24 +11,22 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - """TimesFM models.""" import logging import math import os -from typing import Sequence +from pathlib import Path +from typing import Dict, Optional, Sequence, Union -import huggingface_hub import numpy as np -from safetensors.torch import load_file import torch +from huggingface_hub import ModelHubMixin, hf_hub_download +from safetensors.torch import load_file from torch import nn from .. import configs -from ..torch import dense -from ..torch import transformer -from ..torch import util +from ..torch import dense, transformer, util from . import timesfm_2p5_base revin = util.revin @@ -56,16 +54,11 @@ class TimesFM_2p5_200M_torch_module(nn.Module): # Layers. self.tokenizer = dense.ResidualBlock(self.config.tokenizer) - self.stacked_xf = nn.ModuleList([ - transformer.Transformer(self.config.stacked_transformers.transformer) - for _ in range(self.x) - ]) - self.output_projection_point = dense.ResidualBlock( - self.config.output_projection_point - ) - self.output_projection_quantiles = dense.ResidualBlock( - self.config.output_projection_quantiles + self.stacked_xf = nn.ModuleList( + [transformer.Transformer(self.config.stacked_transformers.transformer) for _ in range(self.x)] ) + self.output_projection_point = dense.ResidualBlock(self.config.output_projection_point) + self.output_projection_quantiles = dense.ResidualBlock(self.config.output_projection_quantiles) # Device. if torch.cuda.is_available(): @@ -82,10 +75,10 @@ class TimesFM_2p5_200M_torch_module(nn.Module): self.to(self.device) def forward( - self, - inputs: torch.Tensor, - masks: torch.Tensor, - decode_caches: list[util.DecodeCache] | None = None, + self, + inputs: torch.Tensor, + masks: torch.Tensor, + decode_caches: list[util.DecodeCache] | None = None, ): tokenizer_inputs = torch.cat([inputs, masks.to(inputs.dtype)], dim=-1) input_embeddings = self.tokenizer(tokenizer_inputs) @@ -96,18 +89,16 @@ class TimesFM_2p5_200M_torch_module(nn.Module): output_embeddings = input_embeddings new_decode_caches = [] for i, layer in enumerate(self.stacked_xf): - output_embeddings, new_cache = layer( - output_embeddings, masks[..., -1], decode_caches[i] - ) + output_embeddings, new_cache = layer(output_embeddings, masks[..., -1], decode_caches[i]) new_decode_caches.append(new_cache) output_ts = self.output_projection_point(output_embeddings) output_quantile_spread = self.output_projection_quantiles(output_embeddings) return ( - input_embeddings, - output_embeddings, - output_ts, - output_quantile_spread, + input_embeddings, + output_embeddings, + output_ts, + output_quantile_spread, ), new_decode_caches def decode(self, horizon: int, inputs, masks): @@ -134,7 +125,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module): patch_sigma = [] for i in range(num_input_patches): (n, mu, sigma), _ = util.update_running_stats( - n, mu, sigma, patched_inputs[:, i], patched_masks[:, i] + n, mu, sigma, patched_inputs[:, i], patched_masks[:, i] ) patch_mu.append(mu) patch_sigma.append(sigma) @@ -143,47 +134,39 @@ class TimesFM_2p5_200M_torch_module(nn.Module): context_sigma = torch.stack(patch_sigma, dim=1) decode_caches = [ - util.DecodeCache( - next_index=torch.zeros( - batch_size, dtype=torch.int32, device=inputs.device - ), - num_masked=torch.zeros( - batch_size, dtype=torch.int32, device=inputs.device - ), - key=torch.zeros( - batch_size, - decode_cache_size, - self.h, - self.hd, - device=inputs.device, - ), - value=torch.zeros( - batch_size, - decode_cache_size, - self.h, - self.hd, - device=inputs.device, - ), - ) - for _ in range(self.x) + util.DecodeCache( + next_index=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device), + num_masked=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device), + key=torch.zeros( + batch_size, + decode_cache_size, + self.h, + self.hd, + device=inputs.device, + ), + value=torch.zeros( + batch_size, + decode_cache_size, + self.h, + self.hd, + device=inputs.device, + ), + ) + for _ in range(self.x) ] - normed_inputs = revin( - patched_inputs, context_mu, context_sigma, reverse=False - ) + normed_inputs = revin(patched_inputs, context_mu, context_sigma, reverse=False) normed_inputs = torch.where(patched_masks, 0.0, normed_inputs) (_, _, normed_outputs, normed_quantile_spread), decode_caches = self( - normed_inputs, patched_masks, decode_caches + normed_inputs, patched_masks, decode_caches ) renormed_outputs = torch.reshape( - revin(normed_outputs, context_mu, context_sigma, reverse=True), - (batch_size, -1, self.o, self.q), + revin(normed_outputs, context_mu, context_sigma, reverse=True), + (batch_size, -1, self.o, self.q), ) renormed_quantile_spread = torch.reshape( - revin( - normed_quantile_spread, context_mu, context_sigma, reverse=True - ), - (batch_size, -1, self.os, self.q), + revin(normed_quantile_spread, context_mu, context_sigma, reverse=True), + (batch_size, -1, self.os, self.q), )[:, -1, ...] # Autogressive decode @@ -191,16 +174,14 @@ class TimesFM_2p5_200M_torch_module(nn.Module): last_renormed_output = renormed_outputs[:, -1, :, self.aridx] for _ in range(num_decode_steps): - new_patched_input = torch.reshape( - last_renormed_output, (batch_size, self.m, self.p) - ) + new_patched_input = torch.reshape(last_renormed_output, (batch_size, self.m, self.p)) new_mask = torch.zeros_like(new_patched_input, dtype=torch.bool) n, mu, sigma = last_n, last_mu, last_sigma new_mus, new_sigmas = [], [] for i in range(self.m): (n, mu, sigma), _ = util.update_running_stats( - n, mu, sigma, new_patched_input[:, i], new_mask[:, i] + n, mu, sigma, new_patched_input[:, i], new_mask[:, i] ) new_mus.append(mu) new_sigmas.append(sigma) @@ -208,16 +189,14 @@ class TimesFM_2p5_200M_torch_module(nn.Module): new_mu = torch.stack(new_mus, dim=1) new_sigma = torch.stack(new_sigmas, dim=1) - new_normed_input = revin( - new_patched_input, new_mu, new_sigma, reverse=False - ) + new_normed_input = revin(new_patched_input, new_mu, new_sigma, reverse=False) (_, _, new_normed_output, _), decode_caches = self( - new_normed_input, new_mask, decode_caches + new_normed_input, new_mask, decode_caches ) new_renormed_output = torch.reshape( - revin(new_normed_output, new_mu, new_sigma, reverse=True), - (batch_size, self.m, self.o, self.q), + revin(new_normed_output, new_mu, new_sigma, reverse=True), + (batch_size, self.m, self.o, self.q), ) ar_outputs.append(new_renormed_output[:, -1, ...]) last_renormed_output = new_renormed_output[:, -1, :, self.aridx] @@ -229,9 +208,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module): return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs - def forecast_naive( - self, horizon: int, inputs: Sequence[np.ndarray] - ) -> list[np.ndarray]: + def forecast_naive(self, horizon: int, inputs: Sequence[np.ndarray]) -> list[np.ndarray]: """Forecasts the time series. This is a naive implementation for debugging purposes. No forecasting @@ -251,12 +228,8 @@ class TimesFM_2p5_200M_torch_module(nn.Module): mask = torch.zeros_like(input_t, dtype=torch.bool) len_front_mask = self.p - (len(each_input) % self.p) if len_front_mask < self.p: - input_t = torch.cat( - [torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0 - ) - mask = torch.cat( - [torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0 - ) + input_t = torch.cat([torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0) + mask = torch.cat([torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0) input_t = input_t[None, ...] mask = mask[None, ...] t_pf, _, t_ar = self.decode(horizon, input_t, mask) @@ -269,38 +242,58 @@ class TimesFM_2p5_200M_torch_module(nn.Module): return outputs -class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): +class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin): """PyTorch implementation of TimesFM 2.5 with 200M parameters.""" model: nn.Module = TimesFM_2p5_200M_torch_module() - def load_checkpoint( - self, - *, - path: str | None = None, - hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch", - ) -> None: - """Loads a PyTorch safetensors TimesFM model. - - Args: - path: Path to a local checkpoint. If not provided, will try to download - from the default Hugging Face repo. - hf_repo_id: If provided, will download from the specified Hugging Face - repo instead. + @classmethod + def _from_pretrained( + cls, + *, + model_id: str, + revision: Optional[str], + cache_dir: Optional[Union[str, Path]], + force_download: bool, + proxies: Optional[Dict], + resume_download: Optional[bool], + local_files_only: bool, + token: Optional[str], + **model_kwargs, + ): """ - if path: - pass - elif hf_repo_id: - logging.info( - "Downloading checkpoint from Hugging Face repo %s", hf_repo_id - ) - path = os.path.join( - huggingface_hub.snapshot_download(hf_repo_id), "model.safetensors" - ) - logging.info("Loading checkpoint from: %s", path) + Loads a PyTorch safetensors TimesFM model from a local path or the Hugging + Face Hub. This method is the backend for the `from_pretrained` class + method provided by `ModelHubMixin`. + """ + # Create an instance of the model wrapper class. + instance = cls(**model_kwargs) + + # Determine the path to the model weights. + model_file_path = "" + if os.path.isdir(model_id): + logging.info("Loading checkpoint from local directory: %s", model_id) + model_file_path = os.path.join(model_id, "model.safetensors") + if not os.path.exists(model_file_path): + raise FileNotFoundError(f"model.safetensors not found in directory {model_id}") else: - raise ValueError("Either path or hf_repo_id must be provided.") - self.model.load_checkpoint(path) + logging.info("Downloading checkpoint from Hugging Face repo %s", model_id) + model_file_path = hf_hub_download( + repo_id=model_id, + filename="model.safetensors", + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + + logging.info("Loading checkpoint from: %s", model_file_path) + # Load the weights into the model. + instance.model.load_checkpoint(model_file_path) + return instance def compile(self, forecast_config: configs.ForecastConfig, **kwargs) -> None: """Attempts to compile the model for fast decoding. @@ -314,49 +307,41 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): if kwargs.get("backend", None) is not None: self.model.compile(**kwargs) - self.global_batch_size = ( - forecast_config.per_core_batch_size * self.model.device_count - ) + self.global_batch_size = forecast_config.per_core_batch_size * self.model.device_count # Shortcut. fc = forecast_config if fc.max_context % self.model.p != 0: logging.info( - "When compiling, max context needs to be multiple of the patch size" - " %d. Using max context = %d instead.", - self.model.p, - new_context := math.ceil(fc.max_context / self.model.p) - * self.model.p, + "When compiling, max context needs to be multiple of the patch size" + " %d. Using max context = %d instead.", + self.model.p, + new_context := math.ceil(fc.max_context / self.model.p) * self.model.p, ) fc.max_context = new_context if fc.max_horizon % self.model.o != 0: logging.info( - "When compiling, max horizon needs to be multiple of the output patch" - " size %d. Using max horizon = %d instead.", - self.model.o, - new_horizon := math.ceil(fc.max_horizon / self.model.o) - * self.model.o, + "When compiling, max horizon needs to be multiple of the output patch" + " size %d. Using max horizon = %d instead.", + self.model.o, + new_horizon := math.ceil(fc.max_horizon / self.model.o) * self.model.o, ) fc.max_horizon = new_horizon if fc.max_context + fc.max_horizon > self.model.config.context_limit: raise ValueError( - "Context + horizon must be less than the context limit." - f" {fc.max_context} + {fc.max_horizon} >" - f" {self.model.config.context_limit}." + "Context + horizon must be less than the context limit." + f" {fc.max_context} + {fc.max_horizon} >" + f" {self.model.config.context_limit}." ) if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os): - raise ValueError( - "Continuous quantile head is not supported for horizons >" - f" {self.model.os}." - ) + raise ValueError(f"Continuous quantile head is not supported for horizons > {self.model.os}.") self.forecast_config = fc def _compiled_decode(horizon, inputs, masks): if horizon > fc.max_horizon: raise ValueError( - "Horizon must be less than the max horizon." - f" {horizon} > {fc.max_horizon}." + f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}." ) inputs = torch.Tensor(np.array(inputs)).to(self.model.device) @@ -376,28 +361,25 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): mu, sigma = None, None pf_outputs, quantile_spreads, ar_outputs = self.model.decode( - forecast_config.max_horizon, inputs, masks + forecast_config.max_horizon, inputs, masks ) to_cat = [pf_outputs[:, -1, ...]] if ar_outputs is not None: to_cat.append(ar_outputs.reshape(batch_size, -1, self.model.q)) full_forecast = torch.cat(to_cat, dim=1) - flip_quantile_fn = lambda x: torch.cat( - [x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1 - ) + def flip_quantile_fn(x): + return torch.cat([x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1) if fc.force_flip_invariance: - flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = ( - self.model.decode(forecast_config.max_horizon, -inputs, masks) + flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = self.model.decode( + forecast_config.max_horizon, -inputs, masks ) flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads) flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs) to_cat = [flipped_pf_outputs[:, -1, ...]] if flipped_ar_outputs is not None: - to_cat.append( - flipped_ar_outputs.reshape(batch_size, -1, self.model.q) - ) + to_cat.append(flipped_ar_outputs.reshape(batch_size, -1, self.model.q)) flipped_full_forecast = torch.cat(to_cat, dim=1) quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2 pf_outputs = (pf_outputs - flipped_pf_outputs) / 2 @@ -406,30 +388,28 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): if fc.use_continuous_quantile_head: for quantile_index in [1, 2, 3, 4, 6, 7, 8, 9]: full_forecast[:, :, quantile_index] = ( - quantile_spreads[:, : fc.max_horizon, quantile_index] - - quantile_spreads[:, : fc.max_horizon, 5] - + full_forecast[:, : fc.max_horizon, 5] + quantile_spreads[:, : fc.max_horizon, quantile_index] + - quantile_spreads[:, : fc.max_horizon, 5] + + full_forecast[:, : fc.max_horizon, 5] ) full_forecast = full_forecast[:, :horizon, :] if fc.return_backcast: - full_backcast = pf_outputs[:, :-1, : self.model.p, :].reshape( - batch_size, -1, self.model.q - ) + full_backcast = pf_outputs[:, :-1, : self.model.p, :].reshape(batch_size, -1, self.model.q) full_forecast = torch.cat([full_backcast, full_forecast], dim=1) if fc.fix_quantile_crossing: for i in [4, 3, 2, 1]: full_forecast[:, :, i] = torch.where( - full_forecast[:, :, i] < full_forecast[:, :, i + 1], - full_forecast[:, :, i], - full_forecast[:, :, i + 1], + full_forecast[:, :, i] < full_forecast[:, :, i + 1], + full_forecast[:, :, i], + full_forecast[:, :, i + 1], ) for i in [6, 7, 8, 9]: full_forecast[:, :, i] = torch.where( - full_forecast[:, :, i] > full_forecast[:, :, i - 1], - full_forecast[:, :, i], - full_forecast[:, :, i - 1], + full_forecast[:, :, i] > full_forecast[:, :, i - 1], + full_forecast[:, :, i], + full_forecast[:, :, i - 1], ) if fc.normalize_inputs: @@ -437,9 +417,9 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): if is_positive is not None: full_forecast = torch.where( - is_positive[..., None], - torch.maximum(full_forecast, torch.zeros_like(full_forecast)), - full_forecast, + is_positive[..., None], + torch.maximum(full_forecast, torch.zeros_like(full_forecast)), + full_forecast, ) full_forecast = full_forecast.detach().cpu().numpy() From 1c6e15ff21c08614d28d496040316c62d739c99a Mon Sep 17 00:00:00 2001 From: Rajat Sen Date: Fri, 26 Sep 2025 22:52:18 +0000 Subject: [PATCH 2/3] update readme --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index f949d8b..cef9d07 100644 --- a/README.md +++ b/README.md @@ -59,8 +59,8 @@ pip install -e . ```python import numpy as np import timesfm -model = timesfm.TimesFM_2p5_200M_torch() -model.load_checkpoint() +model = TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch") + model.compile( timesfm.ForecastConfig( max_context=1024, From 2fded9268710029054b6fdd008046f450dad716b Mon Sep 17 00:00:00 2001 From: Rajat Sen Date: Fri, 26 Sep 2025 23:25:06 +0000 Subject: [PATCH 3/3] indentation changes --- pyproject.toml | 2 +- src/timesfm/timesfm_2p5/timesfm_2p5_torch.py | 45 ++++++++++++++------ 2 files changed, 34 insertions(+), 13 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b11c5d1..734a5b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ safetensors = ">=0.5.3" torch = { version = ">=2.0.0", extras = ["cuda"] } [tool.ruff] -line-length = 100 +line-length = 88 indent-width = 2 diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index e82a0f7..9d201d5 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py @@ -55,10 +55,17 @@ class TimesFM_2p5_200M_torch_module(nn.Module): # Layers. self.tokenizer = dense.ResidualBlock(self.config.tokenizer) self.stacked_xf = nn.ModuleList( - [transformer.Transformer(self.config.stacked_transformers.transformer) for _ in range(self.x)] + [ + transformer.Transformer(self.config.stacked_transformers.transformer) + for _ in range(self.x) + ] + ) + self.output_projection_point = dense.ResidualBlock( + self.config.output_projection_point + ) + self.output_projection_quantiles = dense.ResidualBlock( + self.config.output_projection_quantiles ) - self.output_projection_point = dense.ResidualBlock(self.config.output_projection_point) - self.output_projection_quantiles = dense.ResidualBlock(self.config.output_projection_quantiles) # Device. if torch.cuda.is_available(): @@ -89,7 +96,9 @@ class TimesFM_2p5_200M_torch_module(nn.Module): output_embeddings = input_embeddings new_decode_caches = [] for i, layer in enumerate(self.stacked_xf): - output_embeddings, new_cache = layer(output_embeddings, masks[..., -1], decode_caches[i]) + output_embeddings, new_cache = layer( + output_embeddings, masks[..., -1], decode_caches[i] + ) new_decode_caches.append(new_cache) output_ts = self.output_projection_point(output_embeddings) output_quantile_spread = self.output_projection_quantiles(output_embeddings) @@ -174,7 +183,9 @@ class TimesFM_2p5_200M_torch_module(nn.Module): last_renormed_output = renormed_outputs[:, -1, :, self.aridx] for _ in range(num_decode_steps): - new_patched_input = torch.reshape(last_renormed_output, (batch_size, self.m, self.p)) + new_patched_input = torch.reshape( + last_renormed_output, (batch_size, self.m, self.p) + ) new_mask = torch.zeros_like(new_patched_input, dtype=torch.bool) n, mu, sigma = last_n, last_mu, last_sigma @@ -208,7 +219,9 @@ class TimesFM_2p5_200M_torch_module(nn.Module): return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs - def forecast_naive(self, horizon: int, inputs: Sequence[np.ndarray]) -> list[np.ndarray]: + def forecast_naive( + self, horizon: int, inputs: Sequence[np.ndarray] + ) -> list[np.ndarray]: """Forecasts the time series. This is a naive implementation for debugging purposes. No forecasting @@ -228,7 +241,9 @@ class TimesFM_2p5_200M_torch_module(nn.Module): mask = torch.zeros_like(input_t, dtype=torch.bool) len_front_mask = self.p - (len(each_input) % self.p) if len_front_mask < self.p: - input_t = torch.cat([torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0) + input_t = torch.cat( + [torch.zeros(len_front_mask, dtype=torch.float32), input_t], dim=0 + ) mask = torch.cat([torch.ones(len_front_mask, dtype=torch.bool), mask], dim=0) input_t = input_t[None, ...] mask = mask[None, ...] @@ -307,7 +322,9 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin): if kwargs.get("backend", None) is not None: self.model.compile(**kwargs) - self.global_batch_size = forecast_config.per_core_batch_size * self.model.device_count + self.global_batch_size = ( + forecast_config.per_core_batch_size * self.model.device_count + ) # Shortcut. fc = forecast_config @@ -335,7 +352,9 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin): f" {self.model.config.context_limit}." ) if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os): - raise ValueError(f"Continuous quantile head is not supported for horizons > {self.model.os}.") + raise ValueError( + f"Continuous quantile head is not supported for horizons > {self.model.os}." + ) self.forecast_config = fc def _compiled_decode(horizon, inputs, masks): @@ -372,8 +391,8 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin): return torch.cat([x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1) if fc.force_flip_invariance: - flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = self.model.decode( - forecast_config.max_horizon, -inputs, masks + flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = ( + self.model.decode(forecast_config.max_horizon, -inputs, masks) ) flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads) flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs) @@ -395,7 +414,9 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin): full_forecast = full_forecast[:, :horizon, :] if fc.return_backcast: - full_backcast = pf_outputs[:, :-1, : self.model.p, :].reshape(batch_size, -1, self.model.q) + full_backcast = pf_outputs[:, :-1, : self.model.p, :].reshape( + batch_size, -1, self.model.q + ) full_forecast = torch.cat([full_backcast, full_forecast], dim=1) if fc.fix_quantile_crossing: