Merge pull request #299 from google-research/siriuz42-2.0-pr

MVP readme

README.md

@@ -1,3 +1,84 @@
 # TimesFM
 
-PLACEHOLDER
+TimesFM (Time Series Foundation Model) is a pretrained time-series foundation
+model developed by Google Research for time-series forecasting.
+
+*   Paper:
+    [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688),
+    ICML 2024.
+*   All checkpoints:
+    [TimesFM Hugging Face Collection](https://huggingface.co/collections/google/timesfm-release-66e4be5fdb56e960c1e482a6).
+*   [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/).
+*   [TimesFM in BigQuery](https://cloud.google.com/bigquery/docs/timesfm-model):
+    an official Google product.
+
+This open version is not an officially supported Google product.
+
+**Latest Model Version:** TimesFM 2.5
+
+**Archived Model Versions:**
+
+-   1.0 and 2.0: relevant code archived in the sub directory `v1`. You can `pip
+    install timesfm==1.3.0` to install an older version of this package to load
+    them.
+
+## Update - Sept. 15, 2025
+
+TimesFM 2.5 is out!
+
+Comparing to TimesFM 2.0, this new 2.5 model:
+
+-   uses 200M parameters, down from 500M.
+-   supports up to 16k context length, up from 2048.
+-   supports continuous quantile forecast up to 1k horizon via an optional 30M
+    quantile head.
+-   gets rid of the `frequency` indicator.
+-   has a couple of new forecasting flags.
+
+Along with the model upgrade we have also upgraded the inference API. This repo
+will be under construction over the next few weeks to
+
+1.  add support for an upcoming Flax version of the model (faster inference).
+2.  add back covariate support.
+3.  populate more docstrings, docs and notebook.
+
+### Install
+
+TODO(siriuz42): Package timesfm==2.0.0 and upload to PyPI .
+
+Run
+
+```shell
+git clone https://github.com/google-research/timesfm.git
+cd timesfm
+pip install -e .
+```
+
+### Code Example
+
+```python
+import numpy as np
+import timesfm
+model = timesfm.TimesFM_2p5_200M_torch()
+model.load_checkpoint()
+model.compile(
+    timesfm.ForecastConfig(
+        max_context=1024,
+        max_horizon=256,
+        normalize_inputs=True,
+        use_continuous_quantile_head=True,
+        force_flip_invariance=True,
+        infer_is_positive=True,
+        fix_quantile_crossing=True,
+    )
+)
+point_forecast, quantile_forecast = model.forecast(
+    horizon=12,
+    inputs=[
+        np.linspace(0, 1, 100),
+        np.sin(np.linspace(0, 20, 67)),
+    ],  # Two dummy inputs
+)
+point_forecast.shape  # (2, 12)
+quantile_forecast.shape  # (2, 12, 10): mean, then 10th to 90th quantiles.
+```

src/timesfm/__init__.py

@@ -14,6 +14,7 @@
 
 """TimesFM API."""
 
+from .configs import ForecastConfig
 from .timesfm_2p5 import timesfm_2p5_torch
 
 TimesFM_2p5_200M_torch = timesfm_2p5_torch.TimesFM_2p5_200M_torch

src/timesfm/configs.py

@@ -20,7 +20,33 @@ from typing import Literal
 
 @dataclasses.dataclass(frozen=False)
 class ForecastConfig:
-    """Options for forecasting."""
+  """Options for forecasting.
+
+  Attributes:
+    max_context: The maximum context length. This is used by the complied decode
+      function at inference time during batched inference. Any input time series
+      with length less than max_context will be padded with zeros, and with
+      length greater than max_context will be truncated.
+    max_horizon: The maximum horizon length. This is used by the complied decode
+      function at inference time during batched inference. The compiled cached
+      decoding function will by default forecast till max_horizon.
+    normalize_inputs: Whether to normalize the inputs. This is useful when the
+      raw inputs are of extremely large or small magnitudes which may result in
+      numerical issues.
+    window_size: The window size for decomposed forecasting.
+      TODO(siriuz42):implement it.
+    per_core_batch_size: The batch size per core. Used at inference time during
+      batched inference when multiple GPU / TPU devices are used.
+    use_continuous_quantile_head: Whether to use a separate continuous quantile
+      head to avoid quantile collapsing.
+    force_flip_invariance: Whether to force flip invariance. TimesFM guarantees
+      that TimesFM(aX + b) = a * TimesFM(x) + b for a >= 0 by default. This flag
+      extends it to a < 0 as well.
+    infer_is_positive: Whether to guarantee nonnegativity of the output if the
+      input is nonnegative.
+    fix_quantile_crossing: Whether to fix quantile crossing.
+    return_backcast: Whether to return backcast.
+  """
 
   max_context: int = 0
   max_horizon: int = 0

src/timesfm/timesfm_2p5/timesfm_2p5_base.py

@@ -126,7 +126,13 @@ class TimesFM_2p5_200M_Definition:
 
 
 class TimesFM_2p5:
-    """Abstract base class for TimesFM models."""
+  """Abstract base class for TimesFM models.
+
+  Attributes:
+    forecast_config: Configuration for forecasting flags.
+    compiled_decode: Compiled decode function.
+    global_batch_size: Global batch size.
+  """
 
   forecast_config: ForecastConfig | None = None
   compiled_decode: Callable[..., Any] | None = None

src/timesfm/timesfm_2p5/timesfm_2p5_torch.py

@@ -56,12 +56,10 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
 
     # Layers.
     self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
-        self.stacked_xf = nn.ModuleList(
+    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
     )
@@ -182,7 +180,9 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
           (batch_size, -1, self.o, self.q),
       )
       renormed_quantile_spread = torch.reshape(
-                revin(normed_quantile_spread, context_mu, context_sigma, reverse=True),
+          revin(
+              normed_quantile_spread, context_mu, context_sigma, reverse=True
+          ),
           (batch_size, -1, self.os, self.q),
       )[:, -1, ...]
 
@@ -229,8 +229,22 @@ 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]):
+  def forecast_naive(
-        """Forecasts the time series."""
+      self, horizon: int, inputs: Sequence[np.ndarray]
+  ) -> list[np.ndarray]:
+    """Forecasts the time series.
+
+    This is a naive implementation for debugging purposes. No forecasting
+    flags are used here. Forecasting quality can be subpar.
+
+    Args:
+      horizon: The number of time points to forecast.
+      inputs: A sequence of numpy arrays, each representing a time series to
+        query forecast for.
+
+    Returns:
+      A list of numpy arrays of forecasts.
+    """
     outputs = []
     for each_input in inputs:
       input_t = torch.tensor(each_input, dtype=torch.float32)
@@ -265,12 +279,21 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
       *,
       path: str | None = None,
       hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch",
-    ):
+  ) -> None:
-        """Loads a PyTorch safetensors TimesFM model."""
+    """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.
+    """
     if path:
       pass
     elif hf_repo_id:
-            logging.info("Downloading checkpoint from Hugging Face repo %s", 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"
       )
@@ -279,7 +302,16 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
       raise ValueError("Either path or hf_repo_id must be provided.")
     self.model.load_checkpoint(path)
 
-    def compile(self, forecast_config: configs.ForecastConfig, **kwargs):
+  def compile(self, forecast_config: configs.ForecastConfig, **kwargs) -> None:
+    """Attempts to compile the model for fast decoding.
+
+    See configs.ForecastConfig for more details on the supported flags.
+
+    Args:
+      forecast_config: Configuration for forecasting flags.
+      **kwargs: Additional keyword arguments to pass to model.compile().
+    """
+
     if kwargs.get("backend", None) is not None:
       self.model.compile(**kwargs)
     self.global_batch_size = (
@@ -294,7 +326,8 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
           "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,
+          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:
@@ -302,7 +335,8 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
           "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,
+          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:
@@ -344,13 +378,10 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
       pf_outputs, quantile_spreads, ar_outputs = self.model.decode(
           forecast_config.max_horizon, inputs, masks
       )
-            full_forecast = torch.cat(
+      to_cat = [pf_outputs[:, -1, ...]]
-                [
+      if ar_outputs is not None:
-                    pf_outputs[:, -1, ...],
+        to_cat.append(ar_outputs.reshape(batch_size, -1, self.model.q))
-                    ar_outputs.reshape(batch_size, -1, self.model.q),
+      full_forecast = torch.cat(to_cat, dim=1)
-                ],
-                dim=1,
-            )
 
       flip_quantile_fn = lambda x: torch.cat(
           [x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1
@@ -362,13 +393,12 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
         )
         flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads)
         flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs)
-                flipped_full_forecast = torch.cat(
+        to_cat = [flipped_pf_outputs[:, -1, ...]]
-                    [
+        if flipped_ar_outputs is not None:
-                        flipped_pf_outputs[:, -1, ...],
+          to_cat.append(
-                        flipped_ar_outputs.reshape(batch_size, -1, self.model.q),
+              flipped_ar_outputs.reshape(batch_size, -1, self.model.q)
-                    ],
-                    dim=1,
           )
+        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
         full_forecast = (full_forecast - flipped_full_forecast) / 2

src/timesfm/torch/transformer.py

@@ -86,7 +86,8 @@ class RotaryPositionalEmbedding(nn.Module):
         / self.embedding_dims
     )
     timescale = (
-            self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction
+        self.min_timescale
+        * (self.max_timescale / self.min_timescale) ** fraction
     ).to(inputs.device)
     if position is None:
       seq_length = inputs.shape[1]
@@ -208,16 +209,21 @@ class MultiHeadAttention(nn.Module):
           b, n_patches, dtype=torch.bool, device=inputs_q.device
       )
 
-        query = self.query(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
+    query = self.query(inputs_q).view(
+        b, n_patches, self.num_heads, self.head_dim
+    )
     key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
-        value = self.value(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
+    value = self.value(inputs_q).view(
+        b, n_patches, self.num_heads, self.head_dim
+    )
 
     if decode_cache is None:
       num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1)
       next_index = torch.zeros_like(num_masked, dtype=torch.int32)
     else:
       num_masked = (
-                torch.sum(patch_mask.to(torch.int32), dim=-1) + decode_cache.num_masked
+          torch.sum(patch_mask.to(torch.int32), dim=-1)
+          + decode_cache.num_masked
       )
       next_index = decode_cache.next_index.clone()
 

src/timesfm/torch/util.py

@@ -49,7 +49,9 @@ def update_running_stats(
   inc_mu = inc_mu_numerator / inc_n_safe
   inc_mu = torch.where(inc_n == 0, 0.0, inc_mu)
 
-    inc_var_numerator = torch.sum(((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1)
+  inc_var_numerator = torch.sum(
+      ((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1
+  )
   inc_var = inc_var_numerator / inc_n_safe
   inc_var = torch.where(inc_n == 0, 0.0, inc_var)
   inc_sigma = torch.sqrt(inc_var)