Repository: huggingface/transformers
File: tests/models/timesfm2_5/test_modeling_timesfm2_5.py (Python)

Code before:
            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 test_training(self):
        if not self.model_tester.is_training:
            self.skipTest(reason="ModelTester is not configured to run training tests")

        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_loss_backward(*config_and_inputs)

    def test_loss_computation(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_loss_computation(*config_and_inputs)

    def test_loss_backward(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_loss_backward(*config_and_inputs)

@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,

Review comment (suggestion):
Should be handled by the training tests then so we shouldnt need it; the one above checks a bit more so I'm fine with that test

Code after:
            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))