Merge pull request #398 from darkpowerxo/feat/peft-finetuning-pipeline-2.5

feat: PEFT fine-tuning pipeline (LoRA/DoRA, multi-GPU) for TimesFM 2.5

.github/workflows/main.yml

@@ -10,9 +10,9 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v6
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v6
         with:
           python-version: '3.11'
       - name: Install uv

.github/workflows/manual_publish.yml

@@ -7,9 +7,9 @@ jobs:
   build-and-publish:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v6
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v6
         with:
           python-version: '3.11'
       - name: Install uv

.gitignore

@@ -1,9 +1,11 @@
 .venv/
 dist/
 __pycache__/
+*.egg-info/
 checkpoints/
 wandb/
 datasets/
 results/
-timesfm_jax.egg-info/
+uv.lock
 development_setup.md
+debug.log

README.md

@@ -22,6 +22,12 @@ This open version is not an officially supported Google product.
     install timesfm==1.3.0` to install an older version of this package to load
     them.
 
+## Update - Apr. 9, 2026
+
+Added fine-tuning example using HuggingFace Transformers + PEFT (LoRA) — see
+[`timesfm-forecasting/examples/finetuning/`](timesfm-forecasting/examples/finetuning/).
+Also added unit tests (`tests/`) and incorporated several community fixes.
+
 ## Update - Mar. 19, 2026
 
 Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/tree/master/timesfm-forecasting) is out.
@@ -44,12 +50,13 @@ Comparing to TimesFM 2.0, this new 2.5 model:
 -   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
+Since the Sept. 2025 launch, the following improvements have been completed:
 
-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.
+1.  ✅ Flax version of the model for faster inference.
+2.  ✅ Covariate support via XReg (see Oct. 2025 update).
+3.  ✅ Documentation, examples, and agent skill (see `timesfm-forecasting/`).
+4.  ✅ Fine-tuning example with LoRA via HuggingFace Transformers + PEFT (see `timesfm-forecasting/examples/finetuning/`).
+5.  ✅ Unit tests for core layers, configs, and utilities (see `tests/`).
 
 ### Install
 

src/timesfm/configs.py

@@ -23,11 +23,11 @@ class ForecastConfig:
   """Options for forecasting.
 
   Attributes:
-    max_context: The maximum context length. This is used by the complied decode
+    max_context: The maximum context length. This is used by the compiled 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
+    max_horizon: The maximum horizon length. This is used by the compiled 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

src/timesfm/timesfm_2p5/timesfm_2p5_base.py

@@ -221,9 +221,10 @@ class TimesFM_2p5:
       dynamic_categorical_covariates: A dict of dynamic categorical covariates.
       static_numerical_covariates: A dict of static numerical covariates.
       static_categorical_covariates: A dict of static categorical covariates.
-      xreg_mode: one of "xreg + timesfm" or "timesfm + xreg". "timesfm + xreg"
-        fits a model on the residuals of the TimesFM forecast. "xreg + timesfm"
-        fits a model on the targets then forecasts on the residuals via TimesFM.
+      xreg_mode: one of "xreg + timesfm" or "timesfm + xreg". "xreg + timesfm"
+        first fits an XReg model on the targets, then uses TimesFM to forecast
+        the residuals. "timesfm + xreg" first runs TimesFM to get a forecast,
+        then fits an XReg model on the residuals of that forecast.
       normalize_xreg_target_per_input: whether to normalize the xreg target per
         input in the given batch.
       ridge: ridge penalty for the linear model.
@@ -285,7 +286,7 @@ class TimesFM_2p5:
 
       if test_lens[-1] > self.forecast_config.max_horizon:
         raise ValueError(
-          "Forecast horizon length inferred from the dynamic covaraites is longer than the"
+          "Forecast horizon length inferred from the dynamic covariates is longer than the"
           f"max_horizon defined in the forecast config: {test_lens[-1]} > {self.forecast_config.max_horizon=}."
         )
 

src/timesfm/timesfm_2p5/timesfm_2p5_torch.py

@@ -85,7 +85,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
     if "torch_compile" in kwargs:
       torch_compile = kwargs["torch_compile"]
     if torch_compile:
-      print("Compiling model...")
+      logging.info("Compiling model...")
       self = torch.compile(self)
 
     self.eval()

tests/__init__.py

@@ -0,0 +1 @@
+# Tests for TimesFM.

tests/test_base_utils.py

@@ -0,0 +1,168 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Tests for NaN-handling and interpolation utilities in the base module.
+
+``strip_leading_nans`` and ``linear_interpolation`` sit on the critical
+inference path: every user input passes through them before being
+patched and fed to the transformer.  Incorrect behavior here — silently
+keeping NaN values or interpolating the wrong indices — causes NaN
+propagation through the entire model and produces garbage forecasts.
+"""
+
+import numpy as np
+
+from timesfm.timesfm_2p5.timesfm_2p5_base import (
+  linear_interpolation,
+  strip_leading_nans,
+)
+
+
+# ---------------------------------------------------------------------------
+# strip_leading_nans
+# ---------------------------------------------------------------------------
+
+
+class TestStripLeadingNans:
+  """Tests for strip_leading_nans — removes leading NaN prefix."""
+
+  def test_no_nans_returns_unchanged(self):
+    """An array without NaN values must pass through unmodified."""
+    arr = np.array([1.0, 2.0, 3.0])
+    result = strip_leading_nans(arr)
+    np.testing.assert_array_equal(result, arr)
+
+  def test_strips_leading_nans_only(self):
+    """Leading NaNs are removed; NaNs embedded in the middle are kept."""
+    arr = np.array([np.nan, np.nan, 1.0, np.nan, 3.0])
+    result = strip_leading_nans(arr)
+    expected = np.array([1.0, np.nan, 3.0])
+    np.testing.assert_array_equal(result, expected)
+
+  def test_single_leading_nan(self):
+    """Edge case: exactly one leading NaN."""
+    arr = np.array([np.nan, 5.0, 6.0])
+    result = strip_leading_nans(arr)
+    np.testing.assert_array_equal(result, np.array([5.0, 6.0]))
+
+  def test_no_leading_nan_with_internal_nans(self):
+    """If the first element is valid, nothing is stripped regardless of
+    internal NaNs."""
+    arr = np.array([1.0, np.nan, np.nan, 4.0])
+    result = strip_leading_nans(arr)
+    np.testing.assert_array_equal(result, arr)
+
+  def test_single_valid_element(self):
+    """A single non-NaN element must be returned as-is."""
+    arr = np.array([42.0])
+    result = strip_leading_nans(arr)
+    np.testing.assert_array_equal(result, np.array([42.0]))
+
+  def test_all_nans_returns_full_array(self):
+    """When every element is NaN, ``np.argmax`` on an all-False mask
+    returns 0 — so the implementation returns the original array, not an
+    empty one.
+
+    This documents the *actual* behavior (which differs from the
+    docstring claim of returning an empty array). Downstream code
+    (``linear_interpolation``) is designed to handle this case.
+    """
+    arr = np.array([np.nan, np.nan, np.nan])
+    result = strip_leading_nans(arr)
+    # Actual behavior: argmax(~isnan) = 0 when all NaN → returns full array.
+    assert len(result) == 3
+    assert np.all(np.isnan(result))
+
+  def test_preserves_dtype(self):
+    """Output dtype must match input dtype (float32 stays float32)."""
+    arr = np.array([np.nan, 1.0, 2.0], dtype=np.float32)
+    result = strip_leading_nans(arr)
+    assert result.dtype == np.float32
+
+
+# ---------------------------------------------------------------------------
+# linear_interpolation
+# ---------------------------------------------------------------------------
+
+
+class TestLinearInterpolation:
+  """Tests for linear_interpolation — fills NaN gaps via ``np.interp``."""
+
+  def test_no_nans_returns_identical(self):
+    """Without NaN values the array is returned as-is (fast path)."""
+    arr = np.array([1.0, 2.0, 3.0])
+    result = linear_interpolation(arr.copy())
+    np.testing.assert_array_equal(result, arr)
+
+  def test_interpolates_single_interior_nan(self):
+    """A single interior NaN is linearly interpolated from neighbors."""
+    arr = np.array([0.0, np.nan, 2.0])
+    result = linear_interpolation(arr)
+    np.testing.assert_allclose(result, [0.0, 1.0, 2.0])
+
+  def test_interpolates_multiple_interior_nans(self):
+    """Multiple consecutive interior NaN values are interpolated."""
+    arr = np.array([0.0, np.nan, np.nan, 3.0])
+    result = linear_interpolation(arr)
+    np.testing.assert_allclose(result, [0.0, 1.0, 2.0, 3.0])
+
+  def test_extrapolates_trailing_nans(self):
+    """Trailing NaN values are filled via ``np.interp`` which holds the
+    last known value (nearest-neighbor extrapolation)."""
+    arr = np.array([1.0, 2.0, np.nan, np.nan])
+    result = linear_interpolation(arr)
+    # np.interp extrapolates by clamping to boundary values.
+    np.testing.assert_allclose(result, [1.0, 2.0, 2.0, 2.0])
+
+  def test_extrapolates_leading_nans(self):
+    """Leading NaN values are filled with the first valid value.
+
+    In practice ``strip_leading_nans`` runs first, but the function must
+    still be robust on its own.
+    """
+    arr = np.array([np.nan, np.nan, 3.0, 4.0])
+    result = linear_interpolation(arr)
+    np.testing.assert_allclose(result, [3.0, 3.0, 3.0, 4.0])
+
+  def test_output_has_no_nans(self):
+    """After interpolation, no NaN values should remain."""
+    arr = np.array([np.nan, 1.0, np.nan, np.nan, 4.0, np.nan])
+    result = linear_interpolation(arr)
+    assert not np.any(np.isnan(result))
+
+  def test_preserves_non_nan_values(self):
+    """Non-NaN values in the original array must never be modified."""
+    arr = np.array([10.0, np.nan, 30.0, np.nan, 50.0])
+    original_valid = arr[~np.isnan(arr)].copy()
+    result = linear_interpolation(arr)
+    np.testing.assert_array_equal(
+      result[~np.isnan(np.array([10.0, np.nan, 30.0, np.nan, 50.0]))],
+      original_valid,
+    )
+
+  def test_interpolation_is_monotone_for_monotone_input(self):
+    """If the known values are strictly increasing, the interpolated
+    result must also be non-decreasing — a basic sanity check on the
+    interpolation direction."""
+    arr = np.array([1.0, np.nan, np.nan, 4.0, np.nan, 6.0])
+    result = linear_interpolation(arr)
+    diffs = np.diff(result)
+    assert np.all(diffs >= 0)
+
+  def test_single_non_nan_fills_all_gaps(self):
+    """With only one valid value, every NaN is replaced by that value
+    (np.interp clamps to the single known point)."""
+    arr = np.array([np.nan, 5.0, np.nan])
+    result = linear_interpolation(arr)
+    np.testing.assert_allclose(result, [5.0, 5.0, 5.0])

tests/test_configs.py

@@ -0,0 +1,196 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Tests for TimesFM configuration dataclasses.
+
+These tests verify that config dataclasses enforce immutability, compose
+correctly, and carry the exact default values the model implementation
+relies on. Catching a silent default-value drift here prevents subtle
+inference regressions that would otherwise only surface as degraded
+forecast quality.
+"""
+
+import dataclasses
+
+import pytest
+
+from timesfm.configs import (
+  ForecastConfig,
+  RandomFourierFeaturesConfig,
+  ResidualBlockConfig,
+  StackedTransformersConfig,
+  TransformerConfig,
+)
+
+
+# ---------------------------------------------------------------------------
+# ForecastConfig
+# ---------------------------------------------------------------------------
+
+
+class TestForecastConfig:
+  """Tests for ForecastConfig — the primary user-facing configuration."""
+
+  def test_defaults_match_safe_inference_settings(self):
+    """Default config must be conservative: no normalization, no fancy heads.
+
+    These defaults are what users get when they call ``ForecastConfig()``
+    without arguments. Changing them silently would break all existing
+    code that relies on the defaults.
+    """
+    cfg = ForecastConfig()
+    assert cfg.max_context == 0
+    assert cfg.max_horizon == 0
+    assert cfg.normalize_inputs is False
+    assert cfg.per_core_batch_size == 1
+    assert cfg.use_continuous_quantile_head is False
+    assert cfg.force_flip_invariance is True
+    assert cfg.infer_is_positive is True
+    assert cfg.fix_quantile_crossing is False
+    assert cfg.return_backcast is False
+
+  def test_frozen_prevents_mutation(self):
+    """Configs are frozen dataclasses — mutating them must raise.
+
+    This is critical because ``compile()`` captures the config object and
+    the compiled decode closure relies on its values never changing.
+    """
+    cfg = ForecastConfig(max_context=512)
+    with pytest.raises(dataclasses.FrozenInstanceError):
+      cfg.max_context = 1024
+
+  def test_replace_creates_independent_copy(self):
+    """``dataclasses.replace`` must yield a new object with updated fields.
+
+    The compile path uses ``replace`` to adjust context/horizon to valid
+    multiples; the original config must remain untouched.
+    """
+    original = ForecastConfig(max_context=512, max_horizon=128)
+    replaced = dataclasses.replace(original, max_context=1024)
+
+    assert replaced.max_context == 1024
+    assert replaced.max_horizon == 128  # untouched
+    assert original.max_context == 512  # original unchanged
+
+  def test_equality_is_structural(self):
+    """Two configs with identical fields must be equal (value semantics)."""
+    a = ForecastConfig(max_context=256, normalize_inputs=True)
+    b = ForecastConfig(max_context=256, normalize_inputs=True)
+    assert a == b
+
+  def test_inequality_on_any_field_difference(self):
+    """A single differing field must break equality."""
+    a = ForecastConfig(max_context=256)
+    b = ForecastConfig(max_context=512)
+    assert a != b
+
+
+# ---------------------------------------------------------------------------
+# ResidualBlockConfig
+# ---------------------------------------------------------------------------
+
+
+class TestResidualBlockConfig:
+  """Tests for ResidualBlockConfig used by tokenizer and output projections."""
+
+  def test_frozen_prevents_mutation(self):
+    cfg = ResidualBlockConfig(
+      input_dims=64,
+      hidden_dims=128,
+      output_dims=128,
+      use_bias=True,
+      activation="swish",
+    )
+    with pytest.raises(dataclasses.FrozenInstanceError):
+      cfg.input_dims = 32
+
+  def test_activation_accepts_all_valid_literals(self):
+    """All three activation modes must be constructable without error."""
+    for act in ("relu", "swish", "none"):
+      cfg = ResidualBlockConfig(
+        input_dims=8,
+        hidden_dims=16,
+        output_dims=8,
+        use_bias=False,
+        activation=act,
+      )
+      assert cfg.activation == act
+
+
+# ---------------------------------------------------------------------------
+# TransformerConfig & StackedTransformersConfig
+# ---------------------------------------------------------------------------
+
+
+class TestTransformerConfig:
+  """Tests for TransformerConfig — architecture-level hyperparameters."""
+
+  def test_model_dims_must_be_divisible_by_num_heads(self):
+    """The model instantiation will fail if this invariant is broken.
+
+    We verify the config at least *carries* the right values that the
+    TimesFM 2.5 definition uses (1280 dims, 16 heads → 80 head_dim).
+    """
+    cfg = TransformerConfig(
+      model_dims=1280,
+      hidden_dims=1280,
+      num_heads=16,
+      attention_norm="rms",
+      feedforward_norm="rms",
+      qk_norm="rms",
+      use_bias=False,
+      use_rotary_position_embeddings=True,
+      ff_activation="swish",
+      fuse_qkv=True,
+    )
+    assert cfg.model_dims % cfg.num_heads == 0
+    assert cfg.model_dims // cfg.num_heads == 80  # head_dim
+
+  def test_stacked_config_composes_correctly(self):
+    """StackedTransformersConfig must wrap a TransformerConfig cleanly."""
+    xf = TransformerConfig(
+      model_dims=64,
+      hidden_dims=64,
+      num_heads=4,
+      attention_norm="rms",
+      feedforward_norm="rms",
+      qk_norm="none",
+      use_bias=True,
+      use_rotary_position_embeddings=False,
+      ff_activation="relu",
+      fuse_qkv=False,
+    )
+    stacked = StackedTransformersConfig(num_layers=6, transformer=xf)
+    assert stacked.num_layers == 6
+    assert stacked.transformer is xf
+    assert stacked.transformer.model_dims == 64
+
+
+# ---------------------------------------------------------------------------
+# RandomFourierFeaturesConfig
+# ---------------------------------------------------------------------------
+
+
+class TestRandomFourierFeaturesConfig:
+  """Tests for RandomFourierFeaturesConfig."""
+
+  def test_frozen_prevents_mutation(self):
+    cfg = RandomFourierFeaturesConfig(
+      input_dims=32,
+      output_dims=64,
+      projection_stddev=1.0,
+      use_bias=True,
+    )
+    with pytest.raises(dataclasses.FrozenInstanceError):
+      cfg.output_dims = 128

tests/test_torch_layers.py

@@ -0,0 +1,262 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Tests for PyTorch layer building blocks: ResidualBlock, RMSNorm,
+RandomFourierFeatures.
+
+These layers are the atoms of the TimesFM architecture. Verifying their
+output shapes, numerical properties, and failure modes protects against
+regressions during refactors.  All tests use small dimensions and run
+on CPU — no model checkpoint or GPU required.
+"""
+
+import torch
+import pytest
+
+from timesfm.configs import RandomFourierFeaturesConfig, ResidualBlockConfig
+from timesfm.torch.dense import RandomFourierFeatures, ResidualBlock
+from timesfm.torch.normalization import RMSNorm
+
+
+# ---------------------------------------------------------------------------
+# ResidualBlock
+# ---------------------------------------------------------------------------
+
+
+class TestResidualBlock:
+  """Tests for the residual block: hidden → activation → output + skip."""
+
+  @pytest.fixture
+  def swish_block(self):
+    """A small residual block with SiLU/Swish activation (matches TimesFM)."""
+    cfg = ResidualBlockConfig(
+      input_dims=16,
+      hidden_dims=32,
+      output_dims=8,
+      use_bias=True,
+      activation="swish",
+    )
+    return ResidualBlock(cfg)
+
+  def test_output_shape(self, swish_block):
+    """Output must have the config's ``output_dims`` as the last dimension,
+    regardless of input batch shape."""
+    x = torch.randn(4, 16)
+    out = swish_block(x)
+    assert out.shape == (4, 8)
+
+  def test_output_shape_3d(self, swish_block):
+    """The block must handle (batch, seq, features) inputs — the layout
+    used when processing patched time series."""
+    x = torch.randn(2, 10, 16)
+    out = swish_block(x)
+    assert out.shape == (2, 10, 8)
+
+  def test_residual_connection_nonzero(self):
+    """The residual connection must contribute to the output.
+
+    We verify this by comparing the output when the hidden path is
+    zeroed out vs. the full output.
+    """
+    cfg = ResidualBlockConfig(
+      input_dims=8,
+      hidden_dims=16,
+      output_dims=8,
+      use_bias=False,
+      activation="none",
+    )
+    block = ResidualBlock(cfg)
+    x = torch.randn(2, 8)
+
+    with torch.no_grad():
+      # Residual path only: zero out hidden and output layers.
+      block.hidden_layer.weight.zero_()
+      block.output_layer.weight.zero_()
+      residual_only = block(x)
+
+      # Must equal the residual layer output.
+      expected = block.residual_layer(x)
+      torch.testing.assert_close(residual_only, expected)
+
+  @pytest.mark.parametrize("activation", ["relu", "swish", "none"])
+  def test_all_activations_produce_valid_output(self, activation):
+    """All supported activations must produce finite, non-NaN output."""
+    cfg = ResidualBlockConfig(
+      input_dims=8,
+      hidden_dims=16,
+      output_dims=8,
+      use_bias=True,
+      activation=activation,
+    )
+    block = ResidualBlock(cfg)
+    x = torch.randn(4, 8)
+    out = block(x)
+    assert not torch.any(torch.isnan(out))
+    assert not torch.any(torch.isinf(out))
+
+  def test_invalid_activation_raises(self):
+    """Unsupported activation must raise ``ValueError`` immediately —
+    fail fast rather than producing garbage at inference time."""
+    cfg = ResidualBlockConfig(
+      input_dims=8,
+      hidden_dims=16,
+      output_dims=8,
+      use_bias=True,
+      activation="gelu",
+    )
+    with pytest.raises(ValueError, match="not supported"):
+      ResidualBlock(cfg)
+
+  def test_gradient_flows_through_both_paths(self):
+    """Gradients must reach both the main path and the residual path.
+
+    Dead gradients on either path would prevent the layer from learning.
+    """
+    cfg = ResidualBlockConfig(
+      input_dims=8,
+      hidden_dims=16,
+      output_dims=8,
+      use_bias=True,
+      activation="swish",
+    )
+    block = ResidualBlock(cfg)
+    x = torch.randn(2, 8, requires_grad=True)
+    out = block(x)
+    loss = out.sum()
+    loss.backward()
+
+    assert block.hidden_layer.weight.grad is not None
+    assert block.residual_layer.weight.grad is not None
+    assert torch.any(block.hidden_layer.weight.grad != 0)
+    assert torch.any(block.residual_layer.weight.grad != 0)
+
+
+# ---------------------------------------------------------------------------
+# RMSNorm
+# ---------------------------------------------------------------------------
+
+
+class TestRMSNorm:
+  """Tests for RMS normalization used in transformer attention/FF blocks."""
+
+  def test_output_shape_preserved(self):
+    """RMSNorm must not change the tensor shape."""
+    norm = RMSNorm(num_features=64)
+    x = torch.randn(2, 10, 64)
+    out = norm(x)
+    assert out.shape == x.shape
+
+  def test_zero_scale_produces_zeros(self):
+    """With default scale (initialized to zeros), output must be all zeros.
+
+    This is a critical initialization property: at init, each transformer
+    layer's post-norm effectively passes through zeros, relying on the
+    residual connection to carry signal.
+    """
+    norm = RMSNorm(num_features=8)
+    # scale is initialized to zeros by default.
+    x = torch.randn(4, 8)
+    out = norm(x)
+    torch.testing.assert_close(out, torch.zeros_like(out))
+
+  def test_unit_scale_preserves_rms_magnitude(self):
+    """With scale = 1, output should have approximately unit RMS along
+    the feature dimension — that's the point of RMS normalization."""
+    norm = RMSNorm(num_features=64)
+    with torch.no_grad():
+      norm.scale.fill_(1.0)
+
+    x = torch.randn(8, 64) * 100  # large magnitude
+    out = norm(x)
+
+    rms = torch.sqrt(torch.mean(out**2, dim=-1))
+    # After normalization, RMS should be close to 1.0.
+    torch.testing.assert_close(
+      rms,
+      torch.ones(8),
+      atol=0.1,
+      rtol=0.1,
+    )
+
+  def test_no_nan_on_zero_input(self):
+    """A zero-valued input must not cause NaN (epsilon prevents div-by-0)."""
+    norm = RMSNorm(num_features=8, epsilon=1e-6)
+    with torch.no_grad():
+      norm.scale.fill_(1.0)
+    x = torch.zeros(2, 8)
+    out = norm(x)
+    assert not torch.any(torch.isnan(out))
+
+
+# ---------------------------------------------------------------------------
+# RandomFourierFeatures
+# ---------------------------------------------------------------------------
+
+
+class TestRandomFourierFeatures:
+  """Tests for the random Fourier feature layer."""
+
+  def test_output_shape(self):
+    """Output dims must be exactly ``config.output_dims``."""
+    cfg = RandomFourierFeaturesConfig(
+      input_dims=8,
+      output_dims=32,
+      projection_stddev=1.0,
+      use_bias=True,
+    )
+    layer = RandomFourierFeatures(cfg)
+    x = torch.randn(4, 8)
+    out = layer(x)
+    assert out.shape == (4, 32)
+
+  def test_output_dims_not_multiple_of_4_raises(self):
+    """The four Fourier components (cos, sin, sq_wave_1, sq_wave_2)
+    require ``output_dims`` to be divisible by 4."""
+    cfg = RandomFourierFeaturesConfig(
+      input_dims=8,
+      output_dims=30,  # not divisible by 4
+      projection_stddev=1.0,
+      use_bias=True,
+    )
+    with pytest.raises(ValueError, match="multiple of 4"):
+      RandomFourierFeatures(cfg)
+
+  def test_fourier_components_bounded(self):
+    """cos and sin outputs are bounded in [-1, 1]; sign outputs are
+    bounded in {-1, 0, 1}. The total Fourier part (before residual)
+    is thus bounded. We verify the output stays finite."""
+    cfg = RandomFourierFeaturesConfig(
+      input_dims=8,
+      output_dims=32,
+      projection_stddev=1.0,
+      use_bias=False,
+    )
+    layer = RandomFourierFeatures(cfg)
+    x = torch.randn(16, 8) * 10  # moderately large input
+    out = layer(x)
+    assert not torch.any(torch.isnan(out))
+    assert not torch.any(torch.isinf(out))
+
+  def test_3d_input_supported(self):
+    """The layer must handle (batch, seq, features) tensors."""
+    cfg = RandomFourierFeaturesConfig(
+      input_dims=8,
+      output_dims=16,
+      projection_stddev=1.0,
+      use_bias=True,
+    )
+    layer = RandomFourierFeatures(cfg)
+    x = torch.randn(2, 5, 8)
+    out = layer(x)
+    assert out.shape == (2, 5, 16)

tests/test_torch_utils.py

@@ -0,0 +1,336 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Tests for PyTorch utility functions: running statistics and RevIN.
+
+These utilities are invoked at every patch boundary during autoregressive
+decoding. Bugs here cause silent numerical drift that compounds over
+long horizons, making them especially hard to diagnose from forecast
+output alone.
+"""
+
+import torch
+import numpy as np
+import pytest
+
+from timesfm.torch.util import (
+  DecodeCache,
+  _TOLERANCE,
+  revin,
+  update_running_stats,
+)
+
+
+# ---------------------------------------------------------------------------
+# update_running_stats
+# ---------------------------------------------------------------------------
+
+
+class TestUpdateRunningStats:
+  """Tests for Welford-style online mean / variance accumulation."""
+
+  def test_single_batch_matches_numpy(self):
+    """A single update with no mask must match numpy's mean and std.
+
+    This is the most basic correctness check: feed all values at once
+    and compare against the ground-truth statistics.
+    """
+    x = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
+    mask = torch.zeros_like(x, dtype=torch.bool)
+    n0 = torch.zeros(1)
+    mu0 = torch.zeros(1)
+    sigma0 = torch.zeros(1)
+
+    (new_n, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask)
+
+    np_values = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+    expected_mu = np.mean(np_values)
+    # Population std (ddof=0), same as PyTorch default.
+    expected_sigma = np.std(np_values, ddof=0)
+
+    assert new_n.item() == pytest.approx(5.0)
+    assert new_mu.item() == pytest.approx(expected_mu, abs=1e-5)
+    assert new_sigma.item() == pytest.approx(expected_sigma, abs=1e-5)
+
+  def test_incremental_accumulation_matches_full_computation(self):
+    """Accumulating two batches incrementally must yield the same result
+    as computing statistics over all values at once.
+
+    This is the defining property of online/streaming statistics.
+    """
+    all_values = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+    batch1 = torch.tensor([[1.0, 2.0, 3.0]])
+    batch2 = torch.tensor([[4.0, 5.0, 6.0]])
+
+    def no_mask(t):
+      return torch.zeros_like(t, dtype=torch.bool)
+
+    # Full computation.
+    n0 = torch.zeros(1)
+    mu0 = torch.zeros(1)
+    sigma0 = torch.zeros(1)
+    (full_n, full_mu, full_sigma), _ = update_running_stats(
+      n0, mu0, sigma0, all_values, no_mask(all_values)
+    )
+
+    # Incremental computation.
+    (n1, mu1, sigma1), _ = update_running_stats(
+      n0, mu0, sigma0, batch1, no_mask(batch1)
+    )
+    (inc_n, inc_mu, inc_sigma), _ = update_running_stats(
+      n1, mu1, sigma1, batch2, no_mask(batch2)
+    )
+
+    assert inc_n.item() == pytest.approx(full_n.item())
+    assert inc_mu.item() == pytest.approx(full_mu.item(), abs=1e-5)
+    assert inc_sigma.item() == pytest.approx(full_sigma.item(), abs=1e-5)
+
+  def test_masked_elements_excluded_from_statistics(self):
+    """Masked positions must be completely ignored — as if they don't exist.
+
+    In TimesFM, leading padding is masked. If mask handling is broken,
+    the zero-padding values pollute the running mean and variance.
+    """
+    # Two values: 10 and 20 are valid; 0 is masked.
+    x = torch.tensor([[0.0, 10.0, 20.0]])
+    mask = torch.tensor([[True, False, False]])
+    n0 = torch.zeros(1)
+    mu0 = torch.zeros(1)
+    sigma0 = torch.zeros(1)
+
+    (new_n, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask)
+
+    assert new_n.item() == pytest.approx(2.0)
+    assert new_mu.item() == pytest.approx(15.0, abs=1e-5)
+    expected_sigma = np.std([10.0, 20.0], ddof=0)
+    assert new_sigma.item() == pytest.approx(expected_sigma, abs=1e-5)
+
+  def test_all_masked_yields_zero_stats(self):
+    """When every element is masked, the function must return zeros
+    rather than NaN or raise an error.
+
+    This happens when an input series is entirely padding.
+    """
+    x = torch.tensor([[99.0, 99.0, 99.0]])
+    mask = torch.ones_like(x, dtype=torch.bool)
+    n0 = torch.zeros(1)
+    mu0 = torch.zeros(1)
+    sigma0 = torch.zeros(1)
+
+    (new_n, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask)
+
+    assert new_n.item() == 0.0
+    assert new_mu.item() == 0.0
+    assert new_sigma.item() == 0.0
+
+  def test_batched_computation_independent(self):
+    """Each sample in the batch must be computed independently.
+
+    Cross-sample leakage would corrupt multi-series forecasting.
+    """
+    x = torch.tensor(
+      [
+        [1.0, 2.0, 3.0],
+        [100.0, 200.0, 300.0],
+      ]
+    )
+    mask = torch.zeros_like(x, dtype=torch.bool)
+    n0 = torch.zeros(2)
+    mu0 = torch.zeros(2)
+    sigma0 = torch.zeros(2)
+
+    (new_n, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask)
+
+    assert new_mu[0].item() == pytest.approx(2.0, abs=1e-5)
+    assert new_mu[1].item() == pytest.approx(200.0, abs=1e-5)
+
+    expected_sigma_0 = np.std([1.0, 2.0, 3.0], ddof=0)
+    expected_sigma_1 = np.std([100.0, 200.0, 300.0], ddof=0)
+    assert new_sigma[0].item() == pytest.approx(expected_sigma_0, abs=1e-5)
+    assert new_sigma[1].item() == pytest.approx(expected_sigma_1, abs=1e-5)
+
+  def test_constant_input_yields_zero_sigma(self):
+    """A constant series has zero variance — sigma must be exactly 0.
+
+    This is important because ``revin`` guards against division-by-zero
+    using ``_TOLERANCE`` when sigma is near zero.
+    """
+    x = torch.tensor([[7.0, 7.0, 7.0, 7.0]])
+    mask = torch.zeros_like(x, dtype=torch.bool)
+    n0 = torch.zeros(1)
+    mu0 = torch.zeros(1)
+    sigma0 = torch.zeros(1)
+
+    (_, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask)
+
+    assert new_mu.item() == pytest.approx(7.0)
+    assert new_sigma.item() == pytest.approx(0.0)
+
+
+# ---------------------------------------------------------------------------
+# revin (Reversible Instance Normalization)
+# ---------------------------------------------------------------------------
+
+
+class TestRevIN:
+  """Tests for the RevIN normalization used in patched decoding."""
+
+  def test_forward_then_reverse_is_identity(self):
+    """normalize → denormalize must reconstruct the original tensor.
+
+    This is the fundamental invariant of reversible normalization: the
+    model operates in normalized space, and the output is mapped back
+    to the original scale. Any deviation here directly corrupts the
+    final forecast values.
+    """
+    x = torch.tensor([[10.0, 20.0, 30.0]])
+    mu = torch.tensor([20.0])
+    sigma = torch.tensor([10.0])
+
+    normed = revin(x, mu, sigma, reverse=False)
+    recovered = revin(normed, mu, sigma, reverse=True)
+
+    torch.testing.assert_close(recovered, x, atol=1e-5, rtol=1e-5)
+
+  def test_forward_produces_correct_normalization(self):
+    """After forward normalization: (x - mu) / sigma."""
+    x = torch.tensor([[10.0, 20.0, 30.0]])
+    mu = torch.tensor([20.0])
+    sigma = torch.tensor([10.0])
+
+    normed = revin(x, mu, sigma, reverse=False)
+
+    expected = torch.tensor([[-1.0, 0.0, 1.0]])
+    torch.testing.assert_close(normed, expected, atol=1e-5, rtol=1e-5)
+
+  def test_reverse_produces_correct_denormalization(self):
+    """After reverse: x * sigma + mu."""
+    normed = torch.tensor([[-1.0, 0.0, 1.0]])
+    mu = torch.tensor([20.0])
+    sigma = torch.tensor([10.0])
+
+    recovered = revin(normed, mu, sigma, reverse=True)
+
+    expected = torch.tensor([[10.0, 20.0, 30.0]])
+    torch.testing.assert_close(recovered, expected, atol=1e-5, rtol=1e-5)
+
+  def test_zero_sigma_does_not_produce_nan(self):
+    """When sigma < tolerance, the function substitutes 1.0 to avoid
+    division by zero. This occurs for constant-valued input series.
+
+    NaN propagation from here would poison the entire transformer
+    forward pass.
+    """
+    x = torch.tensor([[5.0, 5.0, 5.0]])
+    mu = torch.tensor([5.0])
+    sigma = torch.tensor([0.0])  # zero variance
+
+    normed = revin(x, mu, sigma, reverse=False)
+
+    assert not torch.any(torch.isnan(normed))
+    assert not torch.any(torch.isinf(normed))
+
+  def test_near_zero_sigma_guarded_by_tolerance(self):
+    """Sigma values just below ``_TOLERANCE`` must trigger the guard."""
+    x = torch.tensor([[1.0, 2.0, 3.0]])
+    mu = torch.tensor([2.0])
+    sigma = torch.tensor([_TOLERANCE / 2])  # below threshold
+
+    normed = revin(x, mu, sigma, reverse=False)
+
+    assert not torch.any(torch.isnan(normed))
+    # With sigma replaced by 1.0: result = x - mu
+    expected = torch.tensor([[-1.0, 0.0, 1.0]])
+    torch.testing.assert_close(normed, expected, atol=1e-5, rtol=1e-5)
+
+  def test_roundtrip_with_batched_3d_input(self):
+    """RevIN must broadcast correctly for (batch, patches, patch_len)
+    tensors — the actual shape used during patched decoding."""
+    batch, patches, patch_len = 2, 4, 32
+    x = torch.randn(batch, patches, patch_len)
+    mu = torch.tensor([1.0, 2.0])  # (batch,)
+    sigma = torch.tensor([3.0, 4.0])  # (batch,)
+
+    normed = revin(x, mu, sigma, reverse=False)
+    recovered = revin(normed, mu, sigma, reverse=True)
+
+    torch.testing.assert_close(recovered, x, atol=1e-5, rtol=1e-5)
+
+  def test_roundtrip_with_batched_4d_input(self):
+    """RevIN must broadcast correctly for (batch, patches, patch_len, q)
+    tensors — the shape used for quantile outputs.
+
+    In the actual decode path, mu/sigma have shape (batch, patches) for
+    4D tensors, so the ``len(mu.shape) == len(x.shape) - 2`` branch
+    fires and adds two trailing singleton dimensions.
+    """
+    batch, patches, patch_len, q = 2, 4, 32, 10
+    x = torch.randn(batch, patches, patch_len, q)
+    # Match the actual call-site shape: (batch, patches)
+    mu = torch.randn(batch, patches)
+    sigma = torch.abs(torch.randn(batch, patches)) + 1.0  # ensure positive
+
+    normed = revin(x, mu, sigma, reverse=False)
+    recovered = revin(normed, mu, sigma, reverse=True)
+
+    torch.testing.assert_close(recovered, x, atol=1e-4, rtol=1e-4)
+
+  def test_negative_values_handled_correctly(self):
+    """RevIN must work for series with negative values (e.g. temperature,
+    financial returns). ``infer_is_positive`` is a separate downstream
+    flag and does not affect RevIN itself.
+    """
+    x = torch.tensor([[-10.0, -5.0, 0.0, 5.0, 10.0]])
+    mu = torch.tensor([0.0])
+    sigma = torch.tensor([7.07])
+
+    normed = revin(x, mu, sigma, reverse=False)
+    recovered = revin(normed, mu, sigma, reverse=True)
+
+    torch.testing.assert_close(recovered, x, atol=1e-3, rtol=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# DecodeCache
+# ---------------------------------------------------------------------------
+
+
+class TestDecodeCache:
+  """Tests for the DecodeCache dataclass used in KV-cache decoding."""
+
+  def test_is_mutable(self):
+    """DecodeCache is *not* frozen — the attention loop mutates
+    ``next_index`` and ``num_masked`` in-place during autoregressive
+    decoding."""
+    cache = DecodeCache(
+      next_index=torch.tensor([0]),
+      num_masked=torch.tensor([0]),
+      key=torch.zeros(1, 10, 4, 8),
+      value=torch.zeros(1, 10, 4, 8),
+    )
+    cache.next_index = torch.tensor([5])
+    assert cache.next_index.item() == 5
+
+  def test_key_value_shape_consistency(self):
+    """Key and value tensors must have identical shapes — they are
+    indexed in parallel during attention computation."""
+    batch, seq, heads, head_dim = 2, 64, 16, 80
+    cache = DecodeCache(
+      next_index=torch.zeros(batch, dtype=torch.int32),
+      num_masked=torch.zeros(batch, dtype=torch.int32),
+      key=torch.zeros(batch, seq, heads, head_dim),
+      value=torch.zeros(batch, seq, heads, head_dim),
+    )
+    assert cache.key.shape == cache.value.shape
+    assert cache.key.shape == (batch, seq, heads, head_dim)

timesfm-forecasting/examples/finetuning/README.md

@@ -0,0 +1,102 @@
+# Fine-Tuning TimesFM 2.5 with LoRA
+
+Parameter-efficient fine-tuning of
+[TimesFM 2.5](https://huggingface.co/google/timesfm-2.5-200m-transformers)
+using **HuggingFace Transformers** and **PEFT (LoRA)**.
+
+This approach is based on the fine-tuning workflow by
+[@kashif](https://github.com/kashif) at HuggingFace
+([notebook](https://github.com/huggingface/notebooks/blob/main/examples/timesfm2_5.ipynb)).
+
+## How It Works
+
+TimesFM 2.5 is available as a standard
+[Transformers](https://github.com/huggingface/transformers) model
+(`TimesFm2_5ModelForPrediction`). This means it supports the full Transformers
+ecosystem out of the box, including:
+
+- **PEFT adapters** — LoRA, QLoRA, etc. via the
+  [`peft`](https://github.com/huggingface/peft) library
+- **All attention backends** — eager, SDPA, Flash Attention 2/3, Flex Attention
+- **Standard `from_pretrained` / `save_pretrained` workflow**
+
+The model's forward pass natively computes a training loss when `future_values`
+are provided, so fine-tuning requires nothing more than a standard PyTorch
+training loop.
+
+## Quick Start
+
+### Install
+
+```bash
+pip install transformers accelerate peft pandas pyarrow scikit-learn
+```
+
+### Train
+
+```bash
+# Fine-tune with default settings on the retail sales dataset
+python finetune_lora.py
+
+# Custom hyperparameters
+python finetune_lora.py \
+    --epochs 20 \
+    --batch_size 64 \
+    --lr 5e-5 \
+    --lora_r 8 \
+    --lora_alpha 16 \
+    --context_len 64 \
+    --horizon_len 13 \
+    --output_dir my-retail-adapter
+```
+
+### Evaluate
+
+```bash
+# Evaluate a previously trained adapter (skip training)
+python finetune_lora.py --eval_only --output_dir timesfm2_5-retail-lora
+```
+
+## Key Concepts
+
+### No External Normalisation
+
+TimesFM 2.5 applies its own internal instance normalisation (RevIN). **Do not**
+normalise your data externally — feed raw values and let the model handle it.
+
+### Random Window Sampling
+
+Following [Chronos-2](https://github.com/amazon-science/chronos-forecasting),
+each training example is a random `(context, horizon)` window sliced from one of
+the input series. This is more data-efficient than always using the same
+fixed window per series.
+
+### LoRA Target Modules
+
+Using `target_modules="all-linear"` applies LoRA to every linear layer in the
+model. With `r=4` this adds only ~0.6% trainable parameters (~1.4M out of
+~232M), which is enough to meaningfully adapt the model to a new domain.
+
+## CLI Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model_id` | `google/timesfm-2.5-200m-transformers` | HuggingFace model ID |
+| `--context_len` | `64` | Context length for training windows |
+| `--horizon_len` | `13` | Forecast horizon in time steps |
+| `--epochs` | `10` | Training epochs |
+| `--batch_size` | `32` | Batch size |
+| `--lr` | `1e-4` | Learning rate |
+| `--lora_r` | `4` | LoRA rank |
+| `--lora_alpha` | `8` | LoRA alpha |
+| `--lora_dropout` | `0.05` | LoRA dropout |
+| `--num_samples` | `5000` | Random training windows to pre-sample |
+| `--output_dir` | `timesfm2_5-retail-lora` | Where to save the adapter |
+| `--seed` | `42` | Random seed |
+| `--eval_only` | — | Skip training; evaluate existing adapter |
+
+## Acknowledgements
+
+The Transformers integration and fine-tuning approach were developed by
+[@kashif](https://github.com/kashif) at HuggingFace. See the original notebook:
+<https://github.com/huggingface/notebooks/blob/main/examples/timesfm2_5.ipynb>

timesfm-forecasting/examples/finetuning/finetune_lora.py

@@ -0,0 +1,446 @@
+#!/usr/bin/env python3
+"""Fine-tune TimesFM 2.5 with LoRA using HuggingFace Transformers + PEFT.
+
+This script demonstrates parameter-efficient fine-tuning of TimesFM 2.5 on a
+retail demand forecasting dataset (weekly store sales).  It uses the HuggingFace
+Transformers checkpoint and the standard PEFT library for LoRA adapters.
+
+The approach is based on the fine-tuning workflow by @kashif at HuggingFace:
+https://github.com/huggingface/notebooks/blob/main/examples/timesfm2_5.ipynb
+
+The dataset is the same one used in the Chronos-2 quickstart notebook.  Each
+store has ~120 weekly data points.  The goal is to forecast the next 13 weeks
+(one quarter) of sales per store.
+
+Requirements:
+    pip install transformers accelerate peft pandas pyarrow scikit-learn
+
+Usage:
+    python finetune_lora.py [OPTIONS]
+
+    Options:
+        --model_id       HuggingFace model ID (default: google/timesfm-2.5-200m-transformers)
+        --context_len    Context length for training windows (default: 64, must be multiple of 32)
+        --horizon_len    Forecast horizon in time steps (default: 13)
+        --epochs         Number of training epochs (default: 10)
+        --batch_size     Training batch size (default: 32)
+        --lr             Learning rate (default: 1e-4)
+        --lora_r         LoRA rank (default: 4)
+        --lora_alpha     LoRA alpha (default: 8)
+        --lora_dropout   LoRA dropout (default: 0.05)
+        --num_samples    Number of random training windows to pre-sample (default: 5000)
+        --output_dir     Directory to save the LoRA adapter (default: timesfm2_5-retail-lora)
+        --seed           Random seed (default: 42)
+"""
+
+import argparse
+import logging
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+class TimeSeriesRandomWindowDataset(Dataset):
+    """Random-window dataset for time series fine-tuning.
+
+    Pre-samples random (series, split-point) windows similar to Chronos-2's
+    random slicing.  Each window has a full *context_len* context (no
+    zero-padding) to avoid corrupting TimesFM's internal RevIN normalisation
+    statistics.
+
+    No external normalisation is needed — TimesFM handles instance
+    normalisation internally.  The loss is computed in the original data scale.
+    """
+
+    def __init__(
+        self,
+        series_list: list[np.ndarray],
+        context_len: int,
+        horizon_len: int,
+        num_samples: int = 5000,
+        seed: int = 42,
+    ):
+        self.series_list = series_list
+        self.context_len = context_len
+        self.horizon_len = horizon_len
+        self.samples: list[tuple[int, int]] = []
+
+        rng = np.random.default_rng(seed)
+        min_len = context_len + horizon_len
+        valid = [i for i, s in enumerate(series_list) if len(s) >= min_len]
+        if not valid:
+            raise ValueError(
+                f"No series long enough for context_len={context_len} + "
+                f"horizon_len={horizon_len}. Shortest series: "
+                f"{min(len(s) for s in series_list)}"
+            )
+
+        for _ in range(num_samples):
+            idx = rng.choice(valid)
+            series = series_list[idx]
+            max_start = len(series) - min_len
+            start = rng.integers(0, max_start + 1)
+            self.samples.append((idx, start))
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, i: int):
+        idx, start = self.samples[i]
+        series = self.series_list[idx]
+        end = start + self.context_len + self.horizon_len
+
+        context = torch.tensor(
+            series[start : start + self.context_len], dtype=torch.float32
+        )
+        target = torch.tensor(
+            series[start + self.context_len : end], dtype=torch.float32
+        )
+        return context, target
+
+
+class TimeSeriesLastWindowDataset(Dataset):
+    """Validation dataset using the last window of each series."""
+
+    def __init__(
+        self,
+        series_list: list[np.ndarray],
+        context_len: int,
+        horizon_len: int,
+    ):
+        self.items: list[tuple[torch.Tensor, torch.Tensor]] = []
+        min_len = context_len + horizon_len
+        for s in series_list:
+            if len(s) >= min_len:
+                ctx = torch.tensor(s[-min_len:-horizon_len], dtype=torch.float32)
+                tgt = torch.tensor(s[-horizon_len:], dtype=torch.float32)
+                self.items.append((ctx, tgt))
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, i: int):
+        return self.items[i]
+
+
+# ---------------------------------------------------------------------------
+# Data loading helpers
+# ---------------------------------------------------------------------------
+
+def load_retail_sales(
+    context_len: int,
+    horizon_len: int,
+    num_samples: int,
+    seed: int,
+) -> tuple[TimeSeriesRandomWindowDataset, TimeSeriesLastWindowDataset]:
+    """Download and prepare the retail sales dataset.
+
+    This is the same dataset used in the Chronos-2 quickstart notebook and
+    in @kashif's TimesFM 2.5 fine-tuning example.  Each store has ~120 weekly
+    data points; the target column is ``Sales``.
+
+    Returns train dataset and val dataset.
+    """
+    logger.info("Loading retail sales dataset …")
+    sales_train_df = pd.read_parquet(
+        "https://autogluon.s3.amazonaws.com/datasets/timeseries/"
+        "retail_sales/train.parquet"
+    )
+    target = "Sales"
+
+    all_series: list[np.ndarray] = []
+    for _, group in sales_train_df.groupby("id"):
+        values = group[target].values.astype(np.float32)
+        if len(values) >= context_len + horizon_len:
+            all_series.append(values)
+
+    logger.info(
+        "Valid stores: %d (need >= %d data points)",
+        len(all_series),
+        context_len + horizon_len,
+    )
+
+    train_ds = TimeSeriesRandomWindowDataset(
+        all_series, context_len, horizon_len, num_samples=num_samples, seed=seed
+    )
+    val_ds = TimeSeriesLastWindowDataset(all_series, context_len, horizon_len)
+    return train_ds, val_ds
+
+
+# ---------------------------------------------------------------------------
+# Training
+# ---------------------------------------------------------------------------
+
+def train(args: argparse.Namespace) -> None:
+    from peft import LoraConfig, get_peft_model
+    from transformers import TimesFm2_5ModelForPrediction
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    logger.info("Using device: %s", device)
+
+    # ------------------------------------------------------------------
+    # Load model
+    # ------------------------------------------------------------------
+    logger.info("Loading model: %s", args.model_id)
+    model = TimesFm2_5ModelForPrediction.from_pretrained(
+        args.model_id,
+        torch_dtype=torch.bfloat16,
+        device_map=device,
+    )
+    horizon_len = args.horizon_len
+    context_len = min(args.context_len, model.config.context_length)
+
+    # ------------------------------------------------------------------
+    # Apply LoRA
+    # ------------------------------------------------------------------
+    lora_config = LoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        target_modules="all-linear",
+        lora_dropout=args.lora_dropout,
+        bias="none",
+    )
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    # ------------------------------------------------------------------
+    # Prepare data
+    # ------------------------------------------------------------------
+    train_ds, val_ds = load_retail_sales(
+        context_len, horizon_len, num_samples=args.num_samples, seed=args.seed
+    )
+    train_loader = DataLoader(
+        train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True
+    )
+    val_loader = DataLoader(val_ds, batch_size=args.batch_size)
+
+    logger.info(
+        "Train samples: %d (%d batches) | Val samples: %d",
+        len(train_ds),
+        len(train_loader),
+        len(val_ds),
+    )
+
+    # ------------------------------------------------------------------
+    # Optimiser & scheduler
+    # ------------------------------------------------------------------
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=args.lr, weight_decay=0.01
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=args.epochs * len(train_loader)
+    )
+
+    # ------------------------------------------------------------------
+    # Training loop
+    # ------------------------------------------------------------------
+    best_val_loss = float("inf")
+
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        epoch_loss = 0.0
+        n_batches = 0
+
+        for context, target_vals in train_loader:
+            context = context.to(device)
+            target_vals = target_vals.to(device)
+
+            outputs = model(
+                past_values=context,
+                future_values=target_vals,
+                forecast_context_len=context_len,
+            )
+            loss = outputs.loss
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            optimizer.step()
+            optimizer.zero_grad()
+            scheduler.step()
+
+            epoch_loss += loss.item()
+            n_batches += 1
+
+        avg_train_loss = epoch_loss / max(n_batches, 1)
+
+        # Validation
+        model.eval()
+        val_loss = 0.0
+        val_batches = 0
+        with torch.no_grad():
+            for context, target_vals in val_loader:
+                context = context.to(device)
+                target_vals = target_vals.to(device)
+                outputs = model(
+                    past_values=context,
+                    future_values=target_vals,
+                    forecast_context_len=context_len,
+                )
+                val_loss += outputs.loss.item()
+                val_batches += 1
+
+        avg_val_loss = val_loss / max(val_batches, 1)
+
+        logger.info(
+            "Epoch %d/%d (%d steps) — train loss: %.4f, val loss: %.4f",
+            epoch,
+            args.epochs,
+            n_batches,
+            avg_train_loss,
+            avg_val_loss,
+        )
+
+        if avg_val_loss < best_val_loss:
+            best_val_loss = avg_val_loss
+            model.save_pretrained(args.output_dir)
+            logger.info("  ✓ saved best adapter → %s", args.output_dir)
+
+    logger.info("Training complete. Best val loss: %.4f", best_val_loss)
+
+
+# ---------------------------------------------------------------------------
+# Evaluation
+# ---------------------------------------------------------------------------
+
+def evaluate(args: argparse.Namespace) -> None:
+    """Compare zero-shot vs fine-tuned on a subset of stores."""
+    from peft import PeftModel
+    from transformers import TimesFm2_5ModelForPrediction
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    logger.info("Loading base model …")
+    base_model = TimesFm2_5ModelForPrediction.from_pretrained(
+        args.model_id,
+        torch_dtype=torch.bfloat16,
+        device_map=device,
+    )
+    base_model.eval()
+    horizon_len = args.horizon_len
+    context_len = min(args.context_len, base_model.config.context_length)
+
+    logger.info("Loading LoRA adapter from %s …", args.output_dir)
+    ft_model = PeftModel.from_pretrained(base_model, args.output_dir)
+    ft_model.eval()
+
+    # --- Load data ---
+    sales_train_df = pd.read_parquet(
+        "https://autogluon.s3.amazonaws.com/datasets/timeseries/"
+        "retail_sales/train.parquet"
+    )
+    sales_test_df = pd.read_parquet(
+        "https://autogluon.s3.amazonaws.com/datasets/timeseries/"
+        "retail_sales/test.parquet"
+    )
+    target = "Sales"
+
+    store_ids = sales_train_df["id"].unique()[:8]
+
+    base_maes: list[float] = []
+    ft_maes: list[float] = []
+
+    for store_id in store_ids:
+        store_train = (
+            sales_train_df[sales_train_df["id"] == store_id][target]
+            .values.astype(np.float32)
+        )
+        store_test = (
+            sales_test_df[sales_test_df["id"] == store_id][target]
+            .values.astype(np.float32)
+        )
+        ground_truth = store_test[:horizon_len]
+        if len(ground_truth) < horizon_len or len(store_train) < context_len:
+            continue
+
+        test_input = torch.tensor(
+            store_train[-context_len:], dtype=torch.float32, device=device
+        ).unsqueeze(0)
+
+        with torch.no_grad():
+            base_out = base_model(past_values=test_input)
+            ft_out = ft_model(past_values=test_input)
+
+        base_forecast = base_out.mean_predictions[0, :horizon_len].float().cpu().numpy()
+        ft_forecast = ft_out.mean_predictions[0, :horizon_len].float().cpu().numpy()
+
+        base_mae = float(np.abs(base_forecast - ground_truth).mean())
+        ft_mae = float(np.abs(ft_forecast - ground_truth).mean())
+        base_maes.append(base_mae)
+        ft_maes.append(ft_mae)
+
+        logger.info(
+            "Store %s — zero-shot MAE: %.2f, LoRA MAE: %.2f",
+            store_id,
+            base_mae,
+            ft_mae,
+        )
+
+    if base_maes:
+        avg_base = np.mean(base_maes)
+        avg_ft = np.mean(ft_maes)
+        improvement = (avg_base - avg_ft) / avg_base * 100
+        logger.info("Average zero-shot MAE: %.2f", avg_base)
+        logger.info("Average LoRA MAE:      %.2f", avg_ft)
+        logger.info("Improvement:           %.1f%%", improvement)
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Fine-tune TimesFM 2.5 with LoRA (Transformers + PEFT)"
+    )
+    p.add_argument(
+        "--model_id",
+        default="google/timesfm-2.5-200m-transformers",
+        help="HuggingFace model ID",
+    )
+    p.add_argument("--context_len", type=int, default=64)
+    p.add_argument("--horizon_len", type=int, default=13)
+    p.add_argument("--epochs", type=int, default=10)
+    p.add_argument("--batch_size", type=int, default=32)
+    p.add_argument("--lr", type=float, default=1e-4)
+    p.add_argument("--lora_r", type=int, default=4)
+    p.add_argument("--lora_alpha", type=int, default=8)
+    p.add_argument("--lora_dropout", type=float, default=0.05)
+    p.add_argument("--num_samples", type=int, default=5000)
+    p.add_argument("--output_dir", default="timesfm2_5-retail-lora")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument(
+        "--eval_only",
+        action="store_true",
+        help="Skip training and only run evaluation",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    if not args.eval_only:
+        train(args)
+
+    if os.path.isdir(args.output_dir):
+        evaluate(args)
+    else:
+        logger.warning(
+            "No adapter found at %s — skipping evaluation.", args.output_dir
+        )
+
+
+if __name__ == "__main__":
+    main()

v1/src/timesfm/data_loader.py

@@ -149,11 +149,12 @@ class TimeSeriesdata(object):
     else:
       epoch_len = self.epoch_len
     for idx in perm[0:epoch_len]:
-      for _ in range(num_ts // self.batch_size + 1):
+      batch_indices = range(0, num_ts, self.batch_size)
+      for batch_idx in batch_indices:
         if self.permute:
           tsidx = np.random.choice(num_ts, size=self.batch_size, replace=False)
         else:
-          tsidx = np.arange(num_ts)
+          tsidx = np.arange(batch_idx, min(batch_idx + self.batch_size, num_ts))
         dtimes = np.arange(idx - hist_len, idx + self.pred_len)
         (
             bts_train,

v1/src/timesfm/xreg_lib.py

@@ -339,12 +339,20 @@ class BatchedInContextXRegBase:
       x_train = np.concatenate(x_train, axis=1)
       x_test = np.concatenate(x_test, axis=1)
 
-      # Normalize for robustness.
-      x_mean = np.mean(x_train, axis=0, keepdims=True)
-      x_std = np.where((w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w,
-                       1.0)
-      x_train = [(x_train - x_mean) / x_std]
-      x_test = [(x_test - x_mean) / x_std]
+      # Normalize per-input for robustness (batch-wide normalization
+      # would make each input's result depend on batch composition).
+      train_splits = np.cumsum(self.train_lens)[:-1]
+      test_splits = np.cumsum(self.test_lens)[:-1]
+      train_parts = np.split(x_train, train_splits, axis=0)
+      test_parts = np.split(x_test, test_splits, axis=0)
+      norm_train, norm_test = [], []
+      for tr, te in zip(train_parts, test_parts):
+        m = np.mean(tr, axis=0, keepdims=True)
+        s = np.where((w := np.std(tr, axis=0, keepdims=True)) > _TOL, w, 1.0)
+        norm_train.append((tr - m) / s)
+        norm_test.append((te - m) / s)
+      x_train = [np.concatenate(norm_train, axis=0)]
+      x_test = [np.concatenate(norm_test, axis=0)]
 
     # Categorical features. Encode one by one.
     one_hot_encoder = preprocessing.OneHotEncoder(
@@ -431,56 +439,55 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase):
         assert_covariate_shapes=assert_covariate_shapes,
     )
 
-    x_train = x_train_raw.copy()
-    if max_rows_per_col:
-      nrows, ncols = x_train.shape
-      if nrows > (w := ncols * max_rows_per_col):
-        subsample = jax.random.choice(
-            jax.random.PRNGKey(max_rows_per_col_sample_seed),
-            nrows,
-            (w,),
-            replace=False,
-        )
-        x_train = x_train[subsample]
-        flat_targets = flat_targets[subsample]
-
     device = jax.devices("cpu")[0] if force_on_cpu else None
-    # Runs jitted version of the solvers which are quicker at the cost of
-    # running jitting during the first time calling. Re-jitting happens whenever
-    # new (padded) shapes are encountered.
-    # Ocassionally it helps with the speed and the accuracy if we force single
-    # thread execution on cpu for accelerator machines:
-    # 1. Avoid moving data to accelarator memory.
-    # 2. Avoid precision loss if any.
-    with jax.default_device(device):
-      x_train_raw = _to_padded_jax_array(x_train_raw)
-      x_train = _to_padded_jax_array(x_train)
-      flat_targets = _to_padded_jax_array(flat_targets)
-      x_test = _to_padded_jax_array(x_test)
-      beta_hat = (jnp.linalg.pinv(
-          x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]),
-          hermitian=True,
-      ) @ x_train.T @ flat_targets)
-      y_hat = x_test @ beta_hat
-      y_hat_context = x_train_raw @ beta_hat if debug_info else None
 
+    # Fit per-input regressions to prevent data leakage across batch items.
     outputs = []
     outputs_context = []
+    train_idx, test_idx = 0, 0
 
-    # Reconstruct the ragged 2-dim batched forecasts from flattened linear fits.
-    train_index, test_index = 0, 0
-    for train_index_delta, test_index_delta in zip(self.train_lens,
-                                                   self.test_lens):
-      outputs.append(np.array(y_hat[test_index:(test_index +
-                                                test_index_delta)]))
-      if debug_info:
-        outputs_context.append(
-            np.array(y_hat_context[train_index:(train_index +
-                                                train_index_delta)]))
-      train_index += train_index_delta
-      test_index += test_index_delta
+    with jax.default_device(device):
+      for trl, tel in zip(self.train_lens, self.test_lens):
+        x_tr = x_train_raw[train_idx : train_idx + trl]
+        x_te = x_test[test_idx : test_idx + tel]
+        y_tr = flat_targets[train_idx : train_idx + trl]
+
+        x_tr_fit = x_tr.copy()
+        if max_rows_per_col:
+          nrows, ncols = x_tr_fit.shape
+          if nrows > (w := ncols * max_rows_per_col):
+            subsample = jax.random.choice(
+                jax.random.PRNGKey(max_rows_per_col_sample_seed),
+                nrows,
+                (w,),
+                replace=False,
+            )
+            x_tr_fit = x_tr_fit[subsample]
+            y_tr = y_tr[subsample]
+
+        x_tr_raw_j = _to_padded_jax_array(x_tr)
+        x_tr_j = _to_padded_jax_array(x_tr_fit)
+        y_tr_j = _to_padded_jax_array(y_tr)
+        x_te_j = _to_padded_jax_array(x_te)
+
+        beta_hat = (jnp.linalg.pinv(
+            x_tr_j.T @ x_tr_j + ridge * jnp.eye(x_tr_j.shape[1]),
+            hermitian=True,
+        ) @ x_tr_j.T @ y_tr_j)
+        outputs.append(np.array(x_te_j @ beta_hat))
+        if debug_info:
+          outputs_context.append(np.array(x_tr_raw_j @ beta_hat))
+
+        train_idx += trl
+        test_idx += tel
 
     if debug_info:
-      return outputs, outputs_context, flat_targets, x_train, x_test
+      return (
+          outputs,
+          outputs_context,
+          _to_padded_jax_array(flat_targets),
+          _to_padded_jax_array(x_train_raw),
+          _to_padded_jax_array(x_test),
+      )
     else:
       return outputs

v1/tests/test_data_loader.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from timesfm.data_loader import TimeSeriesdata
+
+
+def test_train_gen_respects_batch_size_when_permute_is_false(tmp_path: Path) -> None:
+    rows = 12
+    df = pd.DataFrame(
+        {
+            "ds": pd.date_range("2024-01-01", periods=rows, freq="D"),
+            "ts_1": np.arange(rows),
+            "ts_2": np.arange(rows) + 10,
+            "ts_3": np.arange(rows) + 20,
+            "ts_4": np.arange(rows) + 30,
+            "ts_5": np.arange(rows) + 40,
+        }
+    )
+    data_path = tmp_path / "sample.csv"
+    df.to_csv(data_path, index=False)
+
+    loader = TimeSeriesdata(
+        data_path=str(data_path),
+        datetime_col="ds",
+        num_cov_cols=None,
+        cat_cov_cols=None,
+        ts_cols=np.array(["ts_1", "ts_2", "ts_3", "ts_4", "ts_5"]),
+        train_range=[0, 8],
+        val_range=[8, 10],
+        test_range=[10, 12],
+        hist_len=3,
+        pred_len=2,
+        batch_size=2,
+        freq="D",
+        normalize=False,
+        epoch_len=1,
+        holiday=False,
+        permute=False,
+    )
+
+    batches = list(loader.train_gen())
+    ts_indices = [batch[-1].tolist() for batch in batches]
+
+    assert ts_indices == [[0, 1], [2, 3], [4]]
+    for batch in batches:
+        assert len(batch[-1]) <= 2
+        assert batch[0].shape[0] == len(batch[-1])
+        assert batch[3].shape[0] == len(batch[-1])