bug fix, docstring, 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

@@ -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,13 +279,14 @@ 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.
 
     Args:
         path: Path to a local checkpoint. If not provided, will try to download
           from the default Hugging Face repo.
-        hf_repo_id: Use another Hugging Face repo ID.
+        hf_repo_id: If provided, will download from the specified Hugging Face
+          repo instead.
     """
     if path:
       pass
@@ -287,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 = (
@@ -354,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
@@ -372,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