transformers/tests/models/timesfm2_5/test_modeling_timesfm2_5.py

# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# 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.

import random
import unittest

import numpy as np
import torch
from parameterized import parameterized

from transformers import TimesFm2_5Config, is_torch_available
from transformers.testing_utils import require_flash_attn, require_torch, require_torch_accelerator, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION, ModelTesterMixin, floats_tensor


if is_torch_available():
    from transformers import TimesFm2_5ModelForPrediction


class TimesFm2_5ModelTester:
    def __init__(
        self,
        parent,
        patch_length: int = 32,
        context_length: int = 128,
        horizon_length: int = 8,
        num_hidden_layers: int = 1,
        hidden_size: int = 32,  # 2 heads * 16 head_dim
        intermediate_size: int = 64,
        head_dim: int = 16,
        num_heads: int = 2,
        rms_norm_eps: float = 1e-6,
        quantiles: list[float] = [0.1, 0.5, 0.9],
        output_quantile_len: int = 16,
        is_training: bool = True,
        batch_size: int = 2,
    ):
        self.parent = parent
        self.patch_length = patch_length
        self.context_length = context_length
        self.horizon_length = horizon_length
        self.quantiles = quantiles
        self.output_quantile_len = output_quantile_len
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.head_dim = head_dim
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_heads
        self.rms_norm_eps = rms_norm_eps
        self.is_training = is_training
        self.batch_size = batch_size

        # The size of test input
        self.seq_length = context_length // patch_length

    def get_config(self):
        return TimesFm2_5Config(
            patch_length=self.patch_length,
            context_length=self.context_length,
            horizon_length=self.horizon_length,
            quantiles=self.quantiles,
            output_quantile_len=self.output_quantile_len,
            hidden_size=self.hidden_size,
            intermediate_size=self.intermediate_size,
            head_dim=self.head_dim,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            num_key_value_heads=self.num_attention_heads,
            rms_norm_eps=self.rms_norm_eps,
        )

    def get_pipeline_config(self):
        return self.get_config()

    def prepare_config_and_inputs(self):
        forecast_input = torch.stack(
            [
                torch.tensor(np.sin(np.linspace(0, 20, 100)), dtype=torch.float32, device=torch_device),
                torch.tensor(np.cos(np.linspace(0, 20, 100)), dtype=torch.float32, device=torch_device),
            ]
        )
        return self.get_config(), forecast_input

    def prepare_config_and_inputs_for_common(self):
        config, forecast_input = self.prepare_config_and_inputs()
        inputs_dict = {"past_values": forecast_input}
        return config, inputs_dict


@require_torch
class TimesFm2_5ModelTest(ModelTesterMixin, unittest.TestCase):
    all_model_classes = (TimesFm2_5ModelForPrediction,) if is_torch_available() else ()
    test_resize_embeddings = False
    is_encoder_decoder = False
    test_inputs_embeds = False
    test_all_params_have_gradient = False

    def setUp(self):
        self.model_tester = TimesFm2_5ModelTester(self)
        self.config_tester = ConfigTester(self, config_class=TimesFm2_5Config)

    def test_create_and_run_model(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        model = TimesFm2_5ModelForPrediction(config)
        model.to(torch_device)
        model.eval()
        results = model(**inputs_dict)
        assert results.mean_predictions is not None

    @unittest.skip(reason="FA backend not yet supported because of forced masks")
    def test_sdpa_can_dispatch_on_flash(self):
        pass

    @unittest.skip(reason="Model does not have input embeddings")
    def test_model_get_set_embeddings(self):
        pass

    @parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION)
    def test_eager_matches_sdpa_inference(
        self, name, dtype, padding_side, use_attention_mask, output_attentions, enable_kernels
    ):
        """
        TimesFM 2.5 computes its own causal attention mask internally from the input padding,
        so the generic test harness (which injects external attention masks and sets RMSNorm eps=1.0
        on QK-norm layers) is not compatible. This override directly verifies eager vs SDPA equivalence.
        """
        if not self.all_model_classes[0]._supports_sdpa:
            self.skipTest("Model does not support SDPA")

        if dtype == "fp16":
            dtype = torch.float16
        elif dtype == "bf16":
            dtype = torch.bfloat16
        elif dtype == "fp32":
            dtype = torch.float32

        tolerance = {torch.float32: 1e-5, torch.bfloat16: 1e-3, torch.float16: 1e-3}[dtype]

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True

        model_eager = TimesFm2_5ModelForPrediction._from_config(config, attn_implementation="eager")
        model_eager.to(dtype=dtype, device=torch_device)
        model_eager.eval()

        model_sdpa = TimesFm2_5ModelForPrediction._from_config(config, attn_implementation="sdpa")
        model_sdpa.load_state_dict(model_eager.state_dict())
        model_sdpa.to(dtype=dtype, device=torch_device)
        model_sdpa.eval()

        past_values = inputs_dict["past_values"].to(dtype=dtype, device=torch_device)

        with torch.no_grad():
            out_eager = model_eager(past_values=past_values)
            out_sdpa = model_sdpa(past_values=past_values)

        # Compare mean predictions
        self.assertTrue(
            torch.allclose(out_eager.mean_predictions, out_sdpa.mean_predictions, atol=tolerance),
            f"mean_predictions max diff: {(out_eager.mean_predictions - out_sdpa.mean_predictions).abs().max().item():.2e}",
        )
        # Compare full predictions
        self.assertTrue(
            torch.allclose(out_eager.full_predictions, out_sdpa.full_predictions, atol=tolerance),
            f"full_predictions max diff: {(out_eager.full_predictions - out_sdpa.full_predictions).abs().max().item():.2e}",
        )
        # Compare last hidden state
        hs_eager = out_eager.hidden_states[-1]
        hs_sdpa = out_sdpa.hidden_states[-1]
        self.assertTrue(
            torch.allclose(hs_eager, hs_sdpa, atol=tolerance),
            f"hidden_states max diff: {(hs_eager - hs_sdpa).abs().max().item():.2e}",
        )

    def _test_flash_or_flex_attn_inference_equivalence(self, attn_implementation):
        """
        TimesFM 2.5 computes its own attention mask internally, so the generic
        flash/flex equivalence test (which injects external attention masks) does not apply.
        This override directly verifies eager vs flash/flex equivalence.
        """
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        dtype = torch.bfloat16
        tolerance = 1e-2

        model_eager = TimesFm2_5ModelForPrediction._from_config(config, attn_implementation="eager")
        model_eager.to(dtype=dtype, device=torch_device)
        model_eager.eval()

        model_fa = TimesFm2_5ModelForPrediction._from_config(config, attn_implementation=attn_implementation)
        model_fa.load_state_dict(model_eager.state_dict())
        model_fa.to(dtype=dtype, device=torch_device)
        model_fa.eval()

        past_values = inputs_dict["past_values"].to(dtype=dtype, device=torch_device)

        with torch.no_grad():
            out_eager = model_eager(past_values=past_values)
            out_fa = model_fa(past_values=past_values)

        self.assertTrue(
            torch.allclose(out_eager.mean_predictions, out_fa.mean_predictions, atol=tolerance),
            f"mean_predictions max diff: {(out_eager.mean_predictions - out_fa.mean_predictions).abs().max().item():.2e}",
        )
        hs_eager = out_eager.hidden_states[-1]
        hs_fa = out_fa.hidden_states[-1]
        self.assertTrue(
            torch.allclose(hs_eager, hs_fa, atol=tolerance),
            f"hidden_states max diff: {(hs_eager - hs_fa).abs().max().item():.2e}",
        )

    @require_flash_attn
    @require_torch_accelerator
    def test_flash_attn_2_inference_equivalence(self):
        self._test_flash_or_flex_attn_inference_equivalence("flash_attention_2")

    @require_flash_attn
    @require_torch_accelerator
    def test_flash_attn_2_inference_equivalence_right_padding(self):
        self._test_flash_or_flex_attn_inference_equivalence("flash_attention_2")

    def test_retain_grad_hidden_states_attentions(self):
        """
        TimesFM 2.5 specific test for retain_grad since the model returns mean_predictions
        as the first tensor, not last_hidden_state like standard models.
        """
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = self.has_attentions

        # force eager attention to support output attentions
        if self.has_attentions:
            config._attn_implementation = "eager"

        # no need to test all models as different heads yield the same functionality
        model_class = self.all_model_classes[0]
        model = model_class._from_config(config, attn_implementation="eager")
        model.to(torch_device)

        inputs = self._prepare_for_class(inputs_dict, model_class)

        outputs = model(**inputs)

        # TimesFM 2.5 returns mean_predictions as first output, not last_hidden_state
        output_tensor = outputs.mean_predictions

        # Encoder-/Decoder-only models
        if outputs.hidden_states is not None:
            hidden_states = outputs.hidden_states[0]
            hidden_states.retain_grad()

        if self.has_attentions and outputs.attentions is not None:
            attentions = outputs.attentions[0]
            attentions.retain_grad()

        output_tensor.flatten()[0].backward(retain_graph=True)

        if outputs.hidden_states is not None:
            self.assertIsNotNone(hidden_states.grad)

        if self.has_attentions and outputs.attentions is not None:
            self.assertIsNotNone(attentions.grad)

    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
        inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
        if return_labels:
            batch_size = inputs_dict["past_values"].shape[0]
            rng = random.Random(42)
            inputs_dict["future_values"] = floats_tensor([batch_size, self.model_tester.horizon_length], rng=rng)
        return inputs_dict


@require_torch
@slow
class TimesFm2_5ModelIntegrationTests(unittest.TestCase):
    def test_inference(self):
        model = TimesFm2_5ModelForPrediction.from_pretrained(
            "google/timesfm-2.5-200m-transformers", revision="refs/pr/3"
        ).to(torch_device)
        forecast_input = [
            np.sin(np.linspace(0, 20, 100)),
            np.sin(np.linspace(0, 20, 200)),
            np.sin(np.linspace(0, 20, 400)),
        ]
        forecast_input_tensor = [torch.tensor(ts, dtype=torch.float32, device=torch_device) for ts in forecast_input]

        with torch.no_grad():
            output = model(past_values=forecast_input_tensor)

        mean_predictions = output.mean_predictions
        self.assertEqual(mean_predictions.shape, torch.Size([3, model.config.horizon_length]))
        # fmt: off
        expected_slice = torch.tensor(
            [ 0.9745,  1.0047,  0.9707,  0.9161,  0.8041,  0.6829,  0.5378,  0.3563,
              0.1698, -0.0396, -0.2508, -0.4358, -0.6150, -0.7491, -0.8659, -0.9535,
             -1.0024, -0.9977, -0.9557, -0.8840, -0.7716, -0.6092, -0.4526, -0.2582,
             -0.0554,  0.1263,  0.3258,  0.5207,  0.6667,  0.7989,  0.9002,  0.9782,
              0.9848,  0.9877,  0.9339,  0.8473,  0.7109,  0.5525,  0.3799,  0.1756,
             -0.0285, -0.2325, -0.4137, -0.5926, -0.7425, -0.8532, -0.9444, -0.9878,
             -0.9985, -0.9828, -0.8972, -0.7833, -0.6414, -0.4881, -0.2838, -0.0878,
              0.1169,  0.3137,  0.4918,  0.6508,  0.7762,  0.8961,  0.9666,  0.9910
            ],
            device=torch_device)
        # fmt: on
        self.assertTrue(torch.allclose(mean_predictions[0, :64], expected_slice, atol=1e-4))