test: add unit tests for configs, torch layers, utils, and base utils

Apply changes from PR #394 by @cj-wong:
- tests/__init__.py: package marker
- tests/test_base_utils.py: strip_leading_nans + linear_interpolation tests
- tests/test_configs.py: frozen dataclass, defaults, replace, equality tests
- tests/test_torch_layers.py: ResidualBlock, RMSNorm, RandomFourierFeatures
- tests/test_torch_utils.py: update_running_stats, revin, DecodeCache tests

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)