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/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