autogluon.timeseries 1.2.1b20250224__py3-none-any.whl → 1.4.1b20251215__py3-none-any.whl

autogluon/timeseries/models/multi_window/multi_window_model.py

@@ -4,12 +4,13 @@ import logging
 import math
 import os
 import time
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any, Type
 
 import numpy as np
+from typing_extensions import Self
 
 import autogluon.core as ag
-from autogluon.timeseries.dataset.ts_dataframe import TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.models.local.abstract_local_model import AbstractLocalModel
 from autogluon.timeseries.splitter import AbstractWindowSplitter, ExpandingWindowSplitter
@@ -25,10 +26,10 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
     Parameters
     ----------
-    model_base : Union[AbstractTimeSeriesModel, Type[AbstractTimeSeriesModel]]
+    model_base
         The base model to repeatedly train. If a AbstractTimeSeriesModel class, then also provide model_base_kwargs
         which will be used to initialize the model via model_base(**model_base_kwargs).
-    model_base_kwargs : Optional[Dict[str, any]], default = None
+    model_base_kwargs
         kwargs used to initialize model_base if model_base is a class.
     """
 
@@ -37,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):
@@ -57,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
@@ -73,10 +74,6 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def supports_past_covariates(self) -> bool:
         return self.model_base.supports_past_covariates
 
-    @property
-    def supports_cat_covariates(self) -> bool:
-        return self.model_base.supports_cat_covariates
-
     def _get_model_base(self):
         return self.model_base
 
@@ -86,16 +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) -> bool:
+    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[int] = None,
-        val_splitter: AbstractWindowSplitter = None,
-        refit_every_n_windows: Optional[int] = 1,
+        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: AbstractWindowSplitter | None = None,
+        refit_every_n_windows: int | None = 1,
         **kwargs,
     ):
         # TODO: use incremental training for GluonTS models?
@@ -109,13 +109,17 @@ 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: AbstractTimeSeriesModel | None = None
 
         for window_index, (train_fold, val_fold) in enumerate(val_splitter.split(train_data)):
             logger.debug(f"\tWindow {window_index}")
+
             # refit_this_window is always True for the 0th window
             refit_this_window = window_index % refit_every_n_windows == 0
+            assert window_index != 0 or refit_this_window
+
             if time_limit is None:
                 time_left_for_window = None
             else:
@@ -138,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
@@ -148,6 +153,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
             else:
                 time_left_for_prediction = time_limit - (time.time() - global_fit_start_time)
 
+            assert model is not None
             model.score_and_cache_oof(
                 val_fold, store_val_score=True, store_predict_time=True, time_limit=time_left_for_prediction
             )
@@ -172,11 +178,14 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
         # Only the model trained on most recent data is saved & used for prediction
         self.most_recent_model = model
-        self.most_recent_model_folder = most_recent_refit_window
+        assert self.most_recent_model is not None
+
+        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
-        self._oof_predictions = oof_predictions_per_window
-        self.val_score = np.mean([info["val_score"] for info in self.info_per_val_window])
+        self.fit_time = time.time() - global_fit_start_time - self.predict_time  # 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()
@@ -191,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:
@@ -205,27 +214,34 @@ 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
 
-    def get_user_params(self) -> dict:
-        return self.model_base.get_user_params()
+        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()
 
-    def _initialize_covariate_regressor_scaler(self, **kwargs) -> None:
+    def _initialize_transforms_and_regressor(self) -> None:
         # Do not initialize the target_scaler and covariate_regressor in the multi window model!
-        pass
-
-    def initialize(self, **kwargs) -> dict:
-        super().initialize(**kwargs)
-        self.model_base.initialize(**kwargs)
-        return kwargs
+        self.target_scaler = None
+        self.covariate_scaler = None
+        self.covariate_regressor = None
 
     def _get_hpo_train_fn_kwargs(self, **train_fn_kwargs) -> dict:
         train_fn_kwargs["is_bagged_model"] = True
@@ -233,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: str = None, verbose=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)
@@ -244,32 +260,36 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
             most_recent_model.save()
         return save_path
 
-    def persist(self):
+    def persist(self) -> Self:
         if self.most_recent_model is None:
             raise ValueError(f"{self.name} must be fit before persisting")
         self.most_recent_model.persist()
+        return self
 
     @classmethod
     def load(
         cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True
     ) -> AbstractTimeSeriesModel:
         model = super().load(path=path, reset_paths=reset_paths, load_oof=load_oof, verbose=verbose)
-        most_recent_model_path = os.path.join(model.path, model.most_recent_model_folder)
-        model.most_recent_model = model.model_base_type.load(
-            most_recent_model_path,
-            reset_paths=reset_paths,
-            verbose=verbose,
-        )
+        if model.most_recent_model_folder is not None:
+            most_recent_model_path = os.path.join(model.path, model.most_recent_model_folder)
+            model.most_recent_model = model.model_base_type.load(
+                most_recent_model_path,
+                reset_paths=reset_paths,
+                verbose=verbose,
+            )
         return model
 
     def convert_to_refit_full_template(self) -> AbstractTimeSeriesModel:
         # refit_model is an instance of base model type, not MultiWindowBacktestingModel
+        assert self.most_recent_model is not None, "Most recent model is None. Model must be fit first."
         refit_model = self.most_recent_model.convert_to_refit_full_template()
         refit_model.rename(self.name + ag.constants.REFIT_FULL_SUFFIX)
         return refit_model
 
     def convert_to_refit_full_via_copy(self) -> AbstractTimeSeriesModel:
         # refit_model is an instance of base model type, not MultiWindowBacktestingModel
+        assert self.most_recent_model is not None, "Most recent model is None. Model must be fit first."
         refit_model = self.most_recent_model.convert_to_refit_full_via_copy()
         refit_model.rename(self.name + ag.constants.REFIT_FULL_SUFFIX)
         return refit_model

autogluon/timeseries/models/registry.py

@@ -0,0 +1,64 @@
+from abc import ABCMeta
+from dataclasses import dataclass
+from inspect import isabstract
+
+
+@dataclass
+class ModelRecord:
+    model_class: type
+    ag_priority: int
+
+
+class ModelRegistry(ABCMeta):
+    """Registry metaclass for time series models. Ensures that TimeSeriesModel classes
+    which implement this metaclass are automatically registered, in order to centralize
+    access to model types.
+
+    See, https://github.com/faif/python-patterns.
+    """
+
+    REGISTRY: dict[str, ModelRecord] = {}
+
+    def __new__(cls, name, bases, attrs):
+        new_cls = super().__new__(cls, name, bases, attrs)
+
+        if name is not None and not isabstract(new_cls):
+            record = ModelRecord(
+                model_class=new_cls,
+                ag_priority=getattr(new_cls, "ag_priority", 0),
+            )
+            cls._add(name.removesuffix("Model"), record)
+
+            # if the class provides additional aliases, register them too
+            if aliases := attrs.get("ag_model_aliases"):
+                for alias in aliases:
+                    cls._add(alias, record)
+
+        return new_cls
+
+    @classmethod
+    def _add(cls, alias: str, record: ModelRecord) -> None:
+        if alias in cls.REGISTRY:
+            raise ValueError(f"You are trying to define a new model with {alias}, but this model already exists.")
+        cls.REGISTRY[alias] = record
+
+    @classmethod
+    def _get_model_record(cls, alias: str | type) -> ModelRecord:
+        if isinstance(alias, type):
+            alias = alias.__name__
+        alias = alias.removesuffix("Model")
+        if alias not in cls.REGISTRY:
+            raise ValueError(f"Unknown model: {alias}, available models are: {cls.available_aliases()}")
+        return cls.REGISTRY[alias]
+
+    @classmethod
+    def get_model_class(cls, alias: str | type) -> type:
+        return cls._get_model_record(alias).model_class
+
+    @classmethod
+    def get_model_priority(cls, alias: str | type) -> int:
+        return cls._get_model_record(alias).ag_priority
+
+    @classmethod
+    def available_aliases(cls) -> list[str]:
+        return sorted(cls.REGISTRY.keys())

autogluon/timeseries/models/toto/__init__.py

@@ -0,0 +1,3 @@
+from .model import TotoModel
+
+__all__ = ["TotoModel"]

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

@@ -0,0 +1,9 @@
+from .backbone import TotoBackbone
+from .dataset import MaskedTimeseries
+from .forecaster import TotoForecaster
+
+__all__ = [
+    "MaskedTimeseries",
+    "TotoBackbone",
+    "TotoForecaster",
+]

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

@@ -0,0 +1,3 @@
+from .backbone import TotoBackbone
+
+__all__ = ["TotoBackbone"]

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

@@ -0,0 +1,196 @@
+# Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+#
+# This product includes software developed at Datadog (https://www.datadoghq.com/)
+# Copyright 2025 Datadog, Inc.
+
+import logging
+from enum import Enum
+
+import torch
+from einops import rearrange
+from torch.nn.functional import scaled_dot_product_attention
+
+from .rope import TimeAwareRotaryEmbedding
+
+log = logging.getLogger(__name__)
+
+
+class AttentionAxis(Enum):
+    TIME = 1
+    SPACE = 2
+
+
+class BaseMultiheadAttention(torch.nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float,
+        rotary_emb: TimeAwareRotaryEmbedding | None,
+        use_memory_efficient_attention: bool,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads."
+        self.head_dim = embed_dim // num_heads
+        self.rotary_emb = rotary_emb
+
+        # We allocate a single tensor for the q, k, and v projection matrices,
+        # multiply them with the inputs, and then split the projected tensors into q, k, and v using unbind.
+        # This reduces overhead a bit vs. having multiple separate Linear layers,
+        # which need to be initialized, tracked by the optimizer, etc.
+        self.wQKV = torch.nn.Linear(embed_dim, embed_dim * 3)
+        self.dropout = dropout
+        self.use_memory_efficient_attention = use_memory_efficient_attention
+        self.wO = torch.nn.Linear(embed_dim, embed_dim)
+
+        assert not self.use_memory_efficient_attention, (
+            "xformers is not available, so use_memory_efficient_attention must be False"
+        )
+
+        if not hasattr(self, "attention_axis") or self.attention_axis not in (AttentionAxis.TIME, AttentionAxis.SPACE):
+            raise ValueError("Child class must define attention_axis as AttentionAxis.TIME or AttentionAxis.SPACE.")
+
+    def rearrange_inputs(self, inputs: torch.Tensor) -> torch.Tensor:
+        pattern = (
+            "batch variate seq_len embed_dim -> (batch variate) seq_len embed_dim"
+            if self.attention_axis == AttentionAxis.TIME
+            else "batch variate seq_len embed_dim -> (batch seq_len) variate embed_dim"
+        )
+
+        return rearrange(inputs, pattern)
+
+    def get_qkv(
+        self,
+        inputs: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        pattern: str = ""
+        if self.attention_axis == AttentionAxis.TIME and self.use_memory_efficient_attention:
+            pattern = "batch_X_variate seq_len (qkv head_dim n_heads) -> qkv batch_X_variate seq_len n_heads head_dim"
+        elif self.attention_axis == AttentionAxis.TIME and not self.use_memory_efficient_attention:
+            pattern = "batch_X_variate seq_len (qkv head_dim n_heads) -> qkv batch_X_variate n_heads seq_len head_dim"
+        elif self.attention_axis == AttentionAxis.SPACE and self.use_memory_efficient_attention:
+            pattern = "batch_X_seq_len variate (qkv head_dim n_heads) -> qkv batch_X_seq_len variate n_heads head_dim"
+        elif self.attention_axis == AttentionAxis.SPACE and not self.use_memory_efficient_attention:
+            pattern = "batch_X_seq_len variate (qkv head_dim n_heads) -> qkv batch_X_seq_len n_heads variate head_dim"
+
+        assert pattern
+        qkv = self.wQKV(inputs.contiguous())
+        return rearrange(qkv, pattern, qkv=3, head_dim=self.head_dim, n_heads=self.num_heads).unbind(dim=0)
+
+    def positional_embedding(self, q, k, v, kv_cache, layer_idx):
+        # Apply the rotary embeddings
+        seq_pos_offset = 0
+        if self.rotary_emb is not None and self.attention_axis == AttentionAxis.TIME:
+            if kv_cache is not None:
+                seq_pos_offset = kv_cache.seq_len(layer_idx)
+
+            # We need to permute because rotary embeddings expect the sequence dimension to be the second-to-last dimension
+            q, k = self.rotary_emb.rotate_queries_and_keys(q, k, seq_pos_offset=seq_pos_offset)
+
+        if kv_cache is not None and self.attention_axis == AttentionAxis.TIME:
+            # First, we append the current input key and value tensors to the cache.
+            # This concatenates the current key and value tensors to the existing key and value tensors
+            kv_cache.append(layer_idx, (k, v))
+            # Then, we retrieve the key and value tensors from the cache.
+            # This includes all the key and value tensors from previous time steps
+            # as well as the current time step.
+            k, v = kv_cache[layer_idx]
+
+        q = q.contiguous()
+        k = k.contiguous().to(q.dtype)  # Ensure k is the same dtype as q; this is necessary when using mixed precision
+        v = v.contiguous().to(q.dtype)  # Ensure v is the same dtype as q; this is necessary when using mixed precision
+
+        return q, k, v, seq_pos_offset
+
+    def rearrange_output(self, output: torch.Tensor, batch: int, variate: int, seq_len: int) -> torch.Tensor:
+        if self.attention_axis == AttentionAxis.TIME and self.use_memory_efficient_attention:
+            pattern = "(batch variate) seq_len n_heads head_dim -> batch variate seq_len (n_heads head_dim)"
+        elif self.attention_axis == AttentionAxis.TIME and not self.use_memory_efficient_attention:
+            pattern = "(batch variate) n_heads seq_len head_dim -> batch variate seq_len (n_heads head_dim)"
+        elif self.attention_axis == AttentionAxis.SPACE and self.use_memory_efficient_attention:
+            pattern = "(batch seq_len) variate n_heads head_dim -> batch variate seq_len (n_heads head_dim)"
+        elif self.attention_axis == AttentionAxis.SPACE and not self.use_memory_efficient_attention:
+            pattern = "(batch seq_len) n_heads variate head_dim -> batch variate seq_len (n_heads head_dim)"
+
+        return rearrange(output, pattern, batch=batch, variate=variate, seq_len=seq_len)  # type: ignore
+
+    def run_attention(self, attention_mask, q, k, v, seq_pos_offset, dropout, seq_len, variate):
+        # Determine dimension ranges for attention
+        # Ensure the last query vector index is used from the cache
+        q_dim_start, q_dim_end = seq_pos_offset, seq_pos_offset + seq_len
+        kv_dim_start, kv_dim_end = 0, v.shape[1] if self.use_memory_efficient_attention else v.shape[2]
+        if self.attention_axis == AttentionAxis.TIME:
+            attention_mask = (
+                attention_mask[..., q_dim_start:q_dim_end, kv_dim_start:kv_dim_end]
+                if torch.is_tensor(attention_mask)
+                else None
+            )
+            return scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=attention_mask,
+                dropout_p=dropout,
+                is_causal=(attention_mask is None and seq_pos_offset == 0),
+            )
+        elif self.attention_axis == AttentionAxis.SPACE:
+            # We don't use causal masking for space-wise attention
+            attention_mask = (
+                attention_mask[..., kv_dim_start:kv_dim_end, kv_dim_start:kv_dim_end]
+                if torch.is_tensor(attention_mask)
+                else None
+            )
+            return scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, dropout_p=dropout, is_causal=False)
+        else:
+            raise ValueError("Invalid attention axis")
+
+    def forward(
+        self,
+        layer_idx: int,
+        inputs: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        kv_cache=None,
+    ) -> torch.Tensor:
+        batch_size, variate, seq_len, _ = inputs.shape
+        dropout = self.dropout if self.training else 0.0
+
+        rearranged_inputs = self.rearrange_inputs(inputs)
+        q, k, v = self.get_qkv(rearranged_inputs)
+
+        q, k, v, seq_pos_offset = self.positional_embedding(q, k, v, kv_cache, layer_idx)
+
+        output = self.run_attention(attention_mask, q, k, v, seq_pos_offset, dropout, seq_len, variate)
+
+        output = self.rearrange_output(output, batch_size, variate, seq_len)
+        return self.wO(output)
+
+
+class TimeWiseMultiheadAttention(BaseMultiheadAttention):
+    """
+    Computes standard multihead causal attention over the time axis.
+    It does this by flattening out the variates along the batch dimension.
+    It also applies rotary position embeddings to the query and key matrices
+    in order to incorporate relative positional information.
+    """
+
+    attention_axis = AttentionAxis.TIME
+
+
+class SpaceWiseMultiheadAttention(BaseMultiheadAttention):
+    """
+    Computes bidirectional multihead attention over the space axis (i.e. across variates within
+    a multi-variate time series). This is done by flattening out the time axis along the batch dimension.
+    This allows the model to attend to different variates at the same time point. By alternating
+    between time-wise and space-wise attention, the model can learn both temporal and cross-variate
+    dependencies in the data.
+
+    Unlike with time-wise attention, don't apply rotary embeddings here
+    because we want cross-variate attention to be invariant to the order of the variates.
+    """
+
+    attention_axis = AttentionAxis.SPACE
+
+
+MultiHeadAttention = TimeWiseMultiheadAttention | SpaceWiseMultiheadAttention