autogluon.timeseries 1.4.1b20251115__py3-none-any.whl → 1.5.0b20251221__py3-none-any.whl

autogluon/timeseries/models/local/statsforecast.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Optional, Type
+from typing import Any, Type
 
 import numpy as np
 import pandas as pd
@@ -19,7 +19,7 @@ class AbstractStatsForecastModel(AbstractLocalModel):
         local_model_args["season_length"] = seasonal_period
         return local_model_args
 
-    def _get_model_type(self, variant: Optional[str] = None) -> Type:
+    def _get_model_type(self, variant: str | None = None) -> Type:
         raise NotImplementedError
 
     def _get_local_model(self, local_model_args: dict):
@@ -162,7 +162,7 @@ class AutoARIMAModel(AbstractProbabilisticStatsForecastModel):
         local_model_args.setdefault("allowmean", True)
         return local_model_args
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import AutoARIMA
 
         return AutoARIMA
@@ -232,7 +232,7 @@ class ARIMAModel(AbstractProbabilisticStatsForecastModel):
         local_model_args.setdefault("order", (1, 1, 1))
         return local_model_args
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import ARIMA
 
         return ARIMA
@@ -269,7 +269,7 @@ class AutoETSModel(AbstractProbabilisticStatsForecastModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
-    ag_priority = 70
+    ag_priority = 60
     init_time_in_seconds = 0  # C++ models require no compilation
     allowed_local_model_args = [
         "damped",
@@ -277,7 +277,7 @@ class AutoETSModel(AbstractProbabilisticStatsForecastModel):
         "seasonal_period",
     ]
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import AutoETS
 
         return AutoETS
@@ -380,7 +380,7 @@ class DynamicOptimizedThetaModel(AbstractProbabilisticStatsForecastModel):
         "seasonal_period",
     ]
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import DynamicOptimizedTheta
 
         return DynamicOptimizedTheta
@@ -425,7 +425,7 @@ class ThetaModel(AbstractProbabilisticStatsForecastModel):
         "seasonal_period",
     ]
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import Theta
 
         return Theta
@@ -546,7 +546,7 @@ class AutoCESModel(AbstractProbabilisticStatsForecastModel):
         "seasonal_period",
     ]
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import AutoCES
 
         return AutoCES
@@ -610,7 +610,7 @@ class ADIDAModel(AbstractStatsForecastIntermittentDemandModel):
 
     ag_priority = 10
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import ADIDA
 
         return ADIDA
@@ -652,7 +652,7 @@ class CrostonModel(AbstractStatsForecastIntermittentDemandModel):
         "variant",
     ]
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import CrostonClassic, CrostonOptimized, CrostonSBA
 
         model_variants = {
@@ -702,7 +702,7 @@ class IMAPAModel(AbstractStatsForecastIntermittentDemandModel):
 
     ag_priority = 10
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         from statsforecast.models import IMAPA
 
         return IMAPA
@@ -726,7 +726,7 @@ class ZeroModel(AbstractStatsForecastIntermittentDemandModel):
 
     ag_priority = 100
 
-    def _get_model_type(self, variant: Optional[str] = None):
+    def _get_model_type(self, variant: str | None = None):
         # ZeroModel does not depend on a StatsForecast implementation
         raise NotImplementedError
 

autogluon/timeseries/models/multi_window/multi_window_model.py

@@ -4,7 +4,7 @@ import logging
 import math
 import os
 import time
-from typing import Any, Optional, Type, Union
+from typing import Any, Type
 
 import numpy as np
 from typing_extensions import Self
@@ -38,8 +38,8 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
     def __init__(
         self,
-        model_base: Union[AbstractTimeSeriesModel, Type[AbstractTimeSeriesModel]],
-        model_base_kwargs: Optional[dict[str, Any]] = None,
+        model_base: AbstractTimeSeriesModel | Type[AbstractTimeSeriesModel],
+        model_base_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ):
         if inspect.isclass(model_base) and issubclass(model_base, AbstractTimeSeriesModel):
@@ -58,8 +58,8 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
         self.model_base_type = type(self.model_base)
         self.info_per_val_window = []
 
-        self.most_recent_model: Optional[AbstractTimeSeriesModel] = None
-        self.most_recent_model_folder: Optional[str] = None
+        self.most_recent_model: AbstractTimeSeriesModel | None = None
+        self.most_recent_model_folder: str | None = None
         super().__init__(**kwargs)
 
     @property
@@ -83,19 +83,19 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def _is_gpu_available(self) -> bool:
         return self._get_model_base()._is_gpu_available()
 
-    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, Union[int, float]]:
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, int | float]:
         return self._get_model_base().get_minimum_resources(is_gpu_available)
 
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
-        num_cpus: Optional[int] = None,
-        num_gpus: Optional[int] = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int | None = None,
         verbosity: int = 2,
-        val_splitter: Optional[AbstractWindowSplitter] = None,
-        refit_every_n_windows: Optional[int] = 1,
+        val_splitter: AbstractWindowSplitter | None = None,
+        refit_every_n_windows: int | None = 1,
         **kwargs,
     ):
         # TODO: use incremental training for GluonTS models?
@@ -109,9 +109,9 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
         if refit_every_n_windows is None:
             refit_every_n_windows = val_splitter.num_val_windows + 1  # only fit model for the first window
 
-        oof_predictions_per_window = []
+        oof_predictions_per_window: list[TimeSeriesDataFrame] = []
         global_fit_start_time = time.time()
-        model: Optional[AbstractTimeSeriesModel] = None
+        model: AbstractTimeSeriesModel | None = None
 
         for window_index, (train_fold, val_fold) in enumerate(val_splitter.split(train_data)):
             logger.debug(f"\tWindow {window_index}")
@@ -142,6 +142,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
                     train_data=train_fold,
                     val_data=val_fold,
                     time_limit=time_left_for_window,
+                    verbosity=verbosity,
                     **kwargs,
                 )
                 model.fit_time = time.time() - model_fit_start_time
@@ -182,8 +183,9 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
         self.most_recent_model_folder = most_recent_refit_window  # type: ignore
         self.predict_time = self.most_recent_model.predict_time
         self.fit_time = time.time() - global_fit_start_time - self.predict_time  # type: ignore
-        self._oof_predictions = oof_predictions_per_window
-        self.val_score = np.mean([info["val_score"] for info in self.info_per_val_window])  # type: ignore
+        self.cache_oof_predictions(oof_predictions_per_window)
+
+        self.val_score = float(np.mean([info["val_score"] for info in self.info_per_val_window]))
 
     def get_info(self) -> dict:
         info = super().get_info()
@@ -198,7 +200,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         if self.most_recent_model is None:
@@ -212,12 +214,25 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
         store_predict_time: bool = False,
         **predict_kwargs,
     ) -> None:
-        # self.val_score, self.predict_time, self._oof_predictions already saved during _fit()
-        assert self._oof_predictions is not None
-        if store_val_score:
-            assert self.val_score is not None
+        if self._oof_predictions is None or self.most_recent_model is None:
+            raise ValueError(f"{self.name} must be fit before calling score_and_cache_oof")
+
+        # Score on val_data using the most recent model
+        past_data, known_covariates = val_data.get_model_inputs_for_scoring(
+            prediction_length=self.prediction_length, known_covariates_names=self.covariate_metadata.known_covariates
+        )
+        predict_start_time = time.time()
+        val_predictions = self.most_recent_model.predict(
+            past_data, known_covariates=known_covariates, **predict_kwargs
+        )
+
+        self._oof_predictions.append(val_predictions)
+
         if store_predict_time:
-            assert self.predict_time is not None
+            self.predict_time = time.time() - predict_start_time
+
+        if store_val_score:
+            self.val_score = self._score_with_predictions(val_data, val_predictions)
 
     def _get_search_space(self):
         return self.model_base._get_search_space()
@@ -234,7 +249,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
         train_fn_kwargs["init_params"]["model_base_kwargs"] = self.get_params()
         return train_fn_kwargs
 
-    def save(self, path: Optional[str] = None, verbose: bool = True) -> str:
+    def save(self, path: str | None = None, verbose: bool = True) -> str:
         most_recent_model = self.most_recent_model
         self.most_recent_model = None
         save_path = super().save(path, verbose)

autogluon/timeseries/models/registry.py

@@ -1,7 +1,6 @@
 from abc import ABCMeta
 from dataclasses import dataclass
 from inspect import isabstract
-from typing import Union
 
 
 @dataclass
@@ -44,7 +43,7 @@ class ModelRegistry(ABCMeta):
         cls.REGISTRY[alias] = record
 
     @classmethod
-    def _get_model_record(cls, alias: Union[str, type]) -> ModelRecord:
+    def _get_model_record(cls, alias: str | type) -> ModelRecord:
         if isinstance(alias, type):
             alias = alias.__name__
         alias = alias.removesuffix("Model")
@@ -53,11 +52,11 @@ class ModelRegistry(ABCMeta):
         return cls.REGISTRY[alias]
 
     @classmethod
-    def get_model_class(cls, alias: Union[str, type]) -> type:
+    def get_model_class(cls, alias: str | type) -> type:
         return cls._get_model_record(alias).model_class
 
     @classmethod
-    def get_model_priority(cls, alias: Union[str, type]) -> int:
+    def get_model_priority(cls, alias: str | type) -> int:
         return cls._get_model_record(alias).ag_priority
 
     @classmethod

autogluon/timeseries/models/toto/_internal/backbone/attention.py

@@ -5,7 +5,6 @@
 
 import logging
 from enum import Enum
-from typing import Optional, Union
 
 import torch
 from einops import rearrange
@@ -27,7 +26,7 @@ class BaseMultiheadAttention(torch.nn.Module):
         embed_dim: int,
         num_heads: int,
         dropout: float,
-        rotary_emb: Optional[TimeAwareRotaryEmbedding],
+        rotary_emb: TimeAwareRotaryEmbedding | None,
         use_memory_efficient_attention: bool,
     ):
         super().__init__()
@@ -151,7 +150,7 @@ class BaseMultiheadAttention(torch.nn.Module):
         self,
         layer_idx: int,
         inputs: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
+        attention_mask: torch.Tensor | None = None,
         kv_cache=None,
     ) -> torch.Tensor:
         batch_size, variate, seq_len, _ = inputs.shape
@@ -194,4 +193,4 @@ class SpaceWiseMultiheadAttention(BaseMultiheadAttention):
     attention_axis = AttentionAxis.SPACE
 
 
-MultiHeadAttention = Union[TimeWiseMultiheadAttention, SpaceWiseMultiheadAttention]
+MultiHeadAttention = TimeWiseMultiheadAttention | SpaceWiseMultiheadAttention

autogluon/timeseries/models/toto/_internal/backbone/backbone.py

@@ -4,7 +4,7 @@
 # Copyright 2025 Datadog, Inc.
 
 import math
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
 import torch
 
@@ -131,7 +131,7 @@ class TotoBackbone(torch.nn.Module):
         scaler_cls: str,
         output_distribution_classes: list[str],
         spacewise_first: bool = True,
-        output_distribution_kwargs: Optional[dict] = None,
+        output_distribution_kwargs: dict | None = None,
         use_memory_efficient_attention: bool = True,
         stabilize_with_global: bool = True,
         scale_factor_exponent: float = 10.0,
@@ -192,8 +192,8 @@ class TotoBackbone(torch.nn.Module):
         inputs: torch.Tensor,
         input_padding_mask: torch.Tensor,
         id_mask: torch.Tensor,
-        kv_cache: Optional[KVCache] = None,
-        scaling_prefix_length: Optional[int] = None,
+        kv_cache: KVCache | None = None,
+        scaling_prefix_length: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         scaled_inputs: torch.Tensor
         loc: torch.Tensor
@@ -244,8 +244,8 @@ class TotoBackbone(torch.nn.Module):
         inputs: torch.Tensor,
         input_padding_mask: torch.Tensor,
         id_mask: torch.Tensor,
-        kv_cache: Optional[KVCache] = None,
-        scaling_prefix_length: Optional[int] = None,
+        kv_cache: KVCache | None = None,
+        scaling_prefix_length: int | None = None,
     ) -> TotoOutput:
         flattened, loc, scale = self.backbone(
             inputs,

autogluon/timeseries/models/toto/_internal/backbone/rope.py

@@ -3,16 +3,11 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/)
 # Copyright 2025 Datadog, Inc.
 
-from typing import Optional
 
 import torch
 from einops import rearrange
-from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
-from rotary_embedding_torch.rotary_embedding_torch import default
 
-
-def exists(val):
-    return val is not None
+from .rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb, default
 
 
 class TimeAwareRotaryEmbedding(RotaryEmbedding):
@@ -41,8 +36,8 @@ class TimeAwareRotaryEmbedding(RotaryEmbedding):
         self,
         q: torch.Tensor,
         k: torch.Tensor,
-        seq_dim: Optional[int] = None,
-        seq_pos: Optional[torch.Tensor] = None,
+        seq_dim: int | None = None,
+        seq_pos: torch.Tensor | None = None,
         seq_pos_offset: int = 0,
     ):
         """
@@ -78,7 +73,7 @@ class TimeAwareRotaryEmbedding(RotaryEmbedding):
 
         return rotated_q, rotated_k
 
-    def get_scale(self, t: torch.Tensor, seq_len: Optional[int] = None, offset=0):
+    def get_scale(self, t: torch.Tensor, seq_len: int | None = None, offset=0):
         """
         This method is adapted closely from the base class, but it knows how to handle
         when `t` has more than 1 dim (as is the case when we're using time-aware RoPE, and have a different