autogluon.timeseries 1.4.1b20250907__py3-none-any.whl → 1.5.1b20260122__py3-none-any.whl

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

@@ -0,0 +1,262 @@
+# 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 math
+from typing import NamedTuple
+
+import torch
+
+from .distribution import MixtureOfStudentTsOutput
+from .kvcache import KVCache
+from .scaler import CausalPatchStdMeanScaler
+from .transformer import Transformer
+
+
+class TotoOutput(NamedTuple):
+    """
+    Output of the Toto model. Contains the output distribution, the location parameters,
+    and the scale parameters.
+    """
+
+    distribution: torch.distributions.Distribution
+    loc: torch.Tensor
+    scale: torch.Tensor
+
+
+def patchify_id_mask(id_mask: torch.Tensor, patch_size: int) -> torch.Tensor:
+    patched_id_mask = id_mask.unfold(dimension=-1, size=patch_size, step=patch_size)
+    patched_id_mask_min = patched_id_mask.min(-1).values
+    patched_id_mask_max = patched_id_mask.max(-1).values
+    assert torch.eq(patched_id_mask_min, patched_id_mask_max).all(), "Patches cannot span multiple datasets"
+    return patched_id_mask_min
+
+
+class PatchEmbedding(torch.nn.Module):
+    """
+    Multivariate time series patch embedding.
+    Patchifies each variate separately.
+    """
+
+    def __init__(self, patch_size: int, stride: int, embed_dim: int):
+        super().__init__()
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.stride = stride
+        self.projection = torch.nn.Linear(self.patch_size, self.embed_dim)
+
+    def _patchify(self, x: torch.Tensor) -> torch.Tensor:
+        return x.unfold(dimension=-1, size=self.patch_size, step=self.stride)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        id_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert x.shape[-1] % self.patch_size == 0, (
+            f"Series length ({x.shape=}) must be divisible by ({self.patch_size=})"
+        )
+        x_patched: torch.Tensor = self._patchify(x)
+        id_mask_patched: torch.Tensor = self._patchify(id_mask)
+
+        assert torch.eq(id_mask_patched.min(-1).values, id_mask_patched.max(-1).values).all(), (
+            "Patches cannot span multiple datasets"
+        )
+
+        return (
+            self.projection(x_patched),
+            id_mask_patched.min(-1).values,
+        )
+
+
+class TotoBackbone(torch.nn.Module):
+    """
+    Toto (Timeseries-Optimized Transformer for Observability) is a transformer-based model for multivariate
+    time series forecasting. It applies a patch embedding to the input data, followed by a transformer
+    that alternates between time-wise and space-wise attention. The transformer is followed by a linear projection
+    that maps the transformer output to the output distribution.
+
+    The output distribution can be a single distribution (e.g. Gaussian) or a mixture of distributions.
+    If a mixture of distributions is used, the model will learn to predict the mixture weights
+    as well as the parameters of the individual distributions.
+
+    Parameters
+    ----------
+    patch_size
+        Size of the patch to use for the patch embedding.
+    stride
+        Stride to use for the patch embedding.
+    embed_dim
+        Dimension of the model's latent space.
+    num_layers
+        Number of transformer layers to use.
+    num_heads
+        Number of attention heads to use in each self-attention layer.
+    mlp_hidden_dim
+        Dimension of the hidden layer in the feedforward network.
+    dropout
+        Dropout rate to use in the model.
+    spacewise_every_n_layers
+        How many time-wise transformer layers to apply between each space-wise transformer layer.
+    spacewise_first
+        Whether to apply space-wise attention before time-wise attention.
+    scaler_cls
+        Class to use for scaling the input data.
+    output_distribution_classes
+        List of classes to use for the output distribution. If a single class is provided, the model
+        will output a single distribution. If multiple classes are provided, the model will output a
+        learned mixture of distributions.
+    output_distribution_kwargs
+        Keyword arguments to pass to the output distribution class. Note: this currently only works
+        with a single output distribution class.
+    use_memory_efficient_attention:
+        Whether to use memory-efficient attention. If True, the model will use the memory-efficient from xFormers.
+    stabilize_with_global:
+        Whether to use global statistics to stabilize causal statistics by clamping extreme values. Only applies to causal scalers.
+    scale_factor_exponent:
+        Exponent that controls the allowed range of deviation from global scale for causal scalers.
+    """
+
+    def __init__(
+        self,
+        patch_size: int,
+        stride: int,
+        embed_dim: int,
+        num_layers: int,
+        num_heads: int,
+        mlp_hidden_dim: int,
+        dropout: float,
+        spacewise_every_n_layers: int,
+        scaler_cls: str,
+        output_distribution_classes: list[str],
+        spacewise_first: bool = True,
+        output_distribution_kwargs: dict | None = None,
+        use_memory_efficient_attention: bool = True,
+        stabilize_with_global: bool = True,
+        scale_factor_exponent: float = 10.0,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        # strings are used when loading a safetensors checkpoint
+        # Initialize patch-based scalers with the correct patch_size
+
+        self.scaler = CausalPatchStdMeanScaler(
+            patch_size=patch_size,
+            stabilize_with_global=stabilize_with_global,
+            scale_factor_exponent=scale_factor_exponent,
+        )
+        self.patch_embed = PatchEmbedding(patch_size, stride, embed_dim)
+        self.dropout = dropout
+        self.num_layers = num_layers
+        self.use_memory_efficient_attention = use_memory_efficient_attention
+        self.transformer = Transformer(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            num_layers=self.num_layers,
+            mlp_hidden_dim=mlp_hidden_dim,
+            dropout=dropout,
+            spacewise_every_n_layers=spacewise_every_n_layers,
+            spacewise_first=spacewise_first,
+            use_memory_efficient_attention=self.use_memory_efficient_attention,
+        )
+        self.unembed = torch.nn.Linear(embed_dim, embed_dim * patch_size)
+
+        # TODO[BEN] this doesn't need to be a list
+        output_distribution_classes_ = [MixtureOfStudentTsOutput]
+        self.output_distribution = output_distribution_classes_[0](embed_dim, **(output_distribution_kwargs or {}))
+
+    def allocate_kv_cache(
+        self,
+        batch_size: int,
+        num_variates: int,
+        max_time_steps: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> KVCache:
+        return KVCache(
+            batch_size=batch_size,
+            num_variates=num_variates,
+            transformer_layers=list(self.transformer.layers),
+            num_layers=self.num_layers,
+            embed_dim=self.embed_dim,
+            num_heads=self.transformer.layers[0].num_heads,  # type: ignore
+            max_seq_len=math.ceil(max_time_steps / self.patch_embed.stride),
+            device=device,
+            dtype=dtype,
+            use_memory_efficient_attention=self.use_memory_efficient_attention,
+        )
+
+    def backbone(
+        self,
+        inputs: torch.Tensor,
+        input_padding_mask: torch.Tensor,
+        id_mask: torch.Tensor,
+        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
+        scale: torch.Tensor
+
+        # Standard scaling operation, same API but without ID mask.
+        scaled_inputs, loc, scale = self.scaler(
+            inputs,
+            weights=torch.ones_like(inputs, device=inputs.device),
+            padding_mask=input_padding_mask,
+            prefix_length=scaling_prefix_length,
+        )
+
+        if kv_cache is not None:
+            prefix_len = self.patch_embed.stride * kv_cache.current_len(0)
+
+            # Truncate inputs so that the transformer only processes
+            # the last patch in the sequence. We'll use the KVCache
+            # for the earlier patches.
+            scaled_inputs = scaled_inputs[:, :, prefix_len:]
+
+            # As a simplification, when using kv cache we only allow decoding
+            # one step at a time after the initial forward pass.
+            assert (prefix_len == 0) or (scaled_inputs.shape[-1] == self.patch_embed.stride), (
+                "Must decode one step at a time."
+            )
+
+            input_padding_mask = input_padding_mask[:, :, prefix_len:]
+            id_mask = id_mask[:, :, prefix_len:]
+
+        embeddings: torch.Tensor
+        reduced_id_mask: torch.Tensor
+
+        embeddings, reduced_id_mask = self.patch_embed(scaled_inputs, id_mask)
+
+        # Apply the transformer on the embeddings
+        transformed: torch.Tensor = self.transformer(embeddings, reduced_id_mask, kv_cache)
+
+        # Unembed and flatten the sequence
+        unembedded = self.unembed(transformed)
+        batch_size, num_variates, seq_len = unembedded.shape[:3]
+        patch_size = unembedded.shape[-1] // self.embed_dim
+        flattened = unembedded.view(batch_size, num_variates, seq_len * patch_size, self.embed_dim)
+        return flattened, loc, scale
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        input_padding_mask: torch.Tensor,
+        id_mask: torch.Tensor,
+        kv_cache: KVCache | None = None,
+        scaling_prefix_length: int | None = None,
+    ) -> TotoOutput:
+        flattened, loc, scale = self.backbone(
+            inputs,
+            input_padding_mask,
+            id_mask,
+            kv_cache,
+            scaling_prefix_length,
+        )
+
+        return TotoOutput(self.output_distribution(flattened), loc, scale)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device

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

@@ -0,0 +1,70 @@
+# 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.
+
+from abc import ABC
+
+import torch
+import torch.nn.functional as F
+from gluonts.torch.distributions import AffineTransformed
+from gluonts.torch.distributions.studentT import StudentT
+
+
+class DistributionOutput(ABC, torch.nn.Module):
+    pass
+
+
+class StudentTOutput(DistributionOutput):
+    def __init__(self, embed_dim):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.df = torch.nn.Linear(embed_dim, 1)
+        self.loc_proj = torch.nn.Linear(embed_dim, 1)
+        self.scale_proj = torch.nn.Linear(embed_dim, 1)
+
+    def forward(self, inputs, loc=None, scale=None):
+        eps = torch.finfo(inputs.dtype).eps
+        df = 2.0 + F.softplus(self.df(inputs)).clamp_min(eps).squeeze(-1)
+        base_loc = self.loc_proj(inputs).squeeze(-1)
+        base_scale = F.softplus(self.scale_proj(inputs)).clamp_min(eps).squeeze(-1)
+
+        base_dist = torch.distributions.StudentT(df, base_loc, base_scale, validate_args=False)  # type: ignore
+
+        if loc is not None and scale is not None:
+            return AffineTransformed(
+                base_dist,
+                loc=loc,
+                scale=scale,
+            )
+        return base_dist
+
+
+class MixtureOfStudentTsOutput(DistributionOutput):
+    def __init__(
+        self,
+        embed_dim,
+        k_components,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.k_components = k_components
+
+        self.df = torch.nn.Linear(embed_dim, k_components)
+        self.loc_proj = torch.nn.Linear(embed_dim, k_components)
+        self.scale_proj = torch.nn.Linear(embed_dim, k_components)
+        self.mixture_weights = torch.nn.Linear(embed_dim, k_components)
+
+    def forward(self, inputs, loc=None, scale=None):
+        df = 2.0 + F.softplus(self.df(inputs)).clamp_min(torch.finfo(inputs.dtype).eps)
+        loc = self.loc_proj(inputs)
+        scale = F.softplus(self.scale_proj(inputs)).clamp_min(torch.finfo(inputs.dtype).eps)
+        logits = self.mixture_weights(inputs)
+        probs = F.softmax(logits, dim=-1)
+        components = StudentT(df, loc, scale)
+        mixture_distribution = torch.distributions.Categorical(probs=probs)
+
+        return torch.distributions.MixtureSameFamily(
+            mixture_distribution,
+            components,
+        )

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

@@ -0,0 +1,136 @@
+# 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.
+
+from dataclasses import dataclass, field
+
+import torch
+
+from .attention import TimeWiseMultiheadAttention
+
+K = torch.Tensor
+V = torch.Tensor
+KV = tuple[torch.Tensor, torch.Tensor]
+
+
+@dataclass
+class KVCache:
+    """
+    Key/Value cache for storing intermediate attention values
+    during multistep inference. Only stores KV cache for timewise layers, skipping spacewise layers.
+    """
+
+    batch_size: int
+    num_variates: int
+    transformer_layers: list
+    num_layers: int
+    embed_dim: int
+    num_heads: int
+    max_seq_len: int
+    device: torch.device = torch.device("cpu")
+    dtype: torch.dtype = torch.float32
+    use_memory_efficient_attention: bool = True
+
+    _keys: torch.Tensor = field(init=False)
+    _values: torch.Tensor = field(init=False)
+    _current_idx: torch.Tensor = field(init=False)
+    _layer_cache_map: torch.Tensor = field(init=False)
+
+    def __post_init__(self):
+        """
+        - Determine timewise vs. spacewise layers and allocate KV only for timewise.
+        - Create a fast tensor-based mapping from global layer_idx -> timewise layer_idx.
+        """
+        assert self.embed_dim % self.num_heads == 0, "embed_dim must be divisible by num_heads"
+        head_dim = self.embed_dim // self.num_heads
+
+        # Compute which layers are timewise
+        time_layer_indices = [
+            i
+            for i in range(self.num_layers)
+            if isinstance(self.transformer_layers[i].attention, TimeWiseMultiheadAttention)
+        ]
+
+        time_layer_count = max(1, len(time_layer_indices))  # handle edge case for no timewise layers
+        # Allocate for only the timewise layers
+        if self.use_memory_efficient_attention:
+            shape = (
+                time_layer_count,
+                self.batch_size * self.num_variates,
+                self.max_seq_len,
+                self.num_heads,
+                head_dim,
+            )
+        else:
+            shape = (
+                time_layer_count,
+                self.batch_size * self.num_variates,
+                self.num_heads,
+                self.max_seq_len,
+                head_dim,
+            )
+        self._keys = torch.zeros(shape, device=self.device, dtype=self.dtype)
+        self._values = torch.zeros_like(self._keys)
+        self._current_idx = torch.zeros(time_layer_count, device=self.device, dtype=torch.int)
+        # Build a tensor lookup for global -> timewise layer index (default to 0)
+        self._layer_cache_map = torch.zeros((self.num_layers,), dtype=torch.int, device=self.device)
+        for cache_idx, layer_idx in enumerate(time_layer_indices):
+            self._layer_cache_map[layer_idx] = int(cache_idx)  # Assign correct indices
+
+    def __getitem__(self, layer_idx: int) -> KV:
+        cache_idx = int(self._layer_cache_map[layer_idx].item())
+        end_idx = int(self._current_idx[cache_idx].item())
+
+        if self.use_memory_efficient_attention:
+            return self._keys[cache_idx, :, :end_idx, :, :], self._values[cache_idx, :, :end_idx, :, :]
+        else:
+            return self._keys[cache_idx, :, :, :end_idx, :], self._values[cache_idx, :, :, :end_idx, :]
+
+    def current_len(self, cache_idx: int) -> int:
+        return int(self._current_idx[cache_idx].item()) if self._current_idx.numel() > 0 else 0
+
+    def seq_len(self, layer_idx: int) -> int:
+        cache_idx = int(self._layer_cache_map[layer_idx].item())
+        return self.current_len(cache_idx)
+
+    def append(self, layer_idx: int, kv: KV):
+        cache_idx = int(self._layer_cache_map[layer_idx].item())
+        keys, values = kv
+
+        # Validate dimensions
+        assert keys.shape == values.shape, "keys and values must have the same shape"
+        assert keys.shape[0] == self.batch_size * self.num_variates, (
+            "keys and values must have batch_size * num_variates as their first dimension"
+        )
+
+        if self.use_memory_efficient_attention:
+            assert keys.shape[2] == self.num_heads, "keys and values must have num_heads as their third dimension"
+        else:
+            assert keys.shape[1] == self.num_heads, "keys and values must have num_heads as their second dimension"
+        assert keys.shape[3] == self.embed_dim // self.num_heads, (
+            "keys and values must have head_dim as their fourth dimension"
+        )
+
+        start_idx = self._current_idx[cache_idx]
+        if self.use_memory_efficient_attention:
+            end_idx = start_idx + keys.shape[1]
+        else:
+            end_idx = start_idx + keys.shape[2]
+        assert end_idx <= self.max_seq_len, (
+            f"max_seq_len exceeded {end_idx} > {self.max_seq_len}, keys.shape: {keys.shape}"
+        )
+
+        if self.use_memory_efficient_attention:
+            self._keys[cache_idx, :, start_idx:end_idx, :, :] = keys
+            self._values[cache_idx, :, start_idx:end_idx, :, :] = values
+        else:
+            self._keys[cache_idx, :, :, start_idx:end_idx, :] = keys
+            self._values[cache_idx, :, :, start_idx:end_idx, :] = values
+
+        self._current_idx[cache_idx] = end_idx
+
+    def reset(self):
+        self._keys.zero_()
+        self._values.zero_()
+        self._current_idx.zero_()

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

@@ -0,0 +1,89 @@
+# 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 torch
+from einops import rearrange
+
+from .rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb, default
+
+
+class TimeAwareRotaryEmbedding(RotaryEmbedding):
+    """
+    A variant of the rotary position embedding that (optionally) uses the time index
+    to compute the sinusoidal and cosine embeddings. This is useful for
+    time series data, where the time index is the most important positional
+    information.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # If the parent stored `freqs` as a Parameter, remove it and register as a buffer
+        # Register buffer is needed for sharding with FSDP
+        if hasattr(self, "freqs") and isinstance(self.freqs, torch.nn.Parameter):
+            # Extract the underlying Tensor
+            freqs_data = self.freqs.data
+
+            # Remove `freqs` from the module's parameters
+            self._parameters.pop("freqs")
+
+            # Register as non-persistent buffer
+            self.register_buffer("freqs", freqs_data, persistent=False)
+
+    def rotate_queries_and_keys(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        seq_dim: int | None = None,
+        seq_pos: torch.Tensor | None = None,
+        seq_pos_offset: int = 0,
+    ):
+        """
+        This method is the same as the one on the base class, except it allows you to override
+        the sequence position tensor with a custom one. It also removes the ability
+        to cache the position encodings, since we have to compute them dynamically
+        based on the timesteps in the input data.
+        """
+        if seq_dim is None:
+            seq_dim = self.default_seq_dim
+
+        assert self.use_xpos
+        device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim]
+
+        seq = default(seq_pos, self.get_seq_pos(seq_len, dtype=dtype, device=device))
+        seq = seq + seq_pos_offset  # type: ignore
+
+        freqs = self.forward(seq)
+
+        scale = self.get_scale(seq).to(dtype)
+
+        # used for xformers
+        if seq_dim == -3:
+            num_heads = q.shape[-2]
+            freqs = freqs.unsqueeze(1).expand(-1, num_heads, -1)
+            scale = scale.unsqueeze(1).expand(-1, num_heads, -1)
+
+        rotated_q = apply_rotary_emb(freqs, q, scale=scale, seq_dim=seq_dim)  # type: ignore
+        rotated_k = apply_rotary_emb(freqs, k, scale=scale**-1, seq_dim=seq_dim)  # type: ignore
+
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+
+        return rotated_q, rotated_k
+
+    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
+        sequence position vector for each time series).
+        """
+        assert self.use_xpos
+
+        power = (t - t.max(-1).values.unsqueeze(-1) // 2) / self.scale_base
+
+        scale = self.scale ** rearrange(power, "... n -> ... n 1")  # type: ignore
+        scale = torch.cat((scale, scale), dim=-1)
+
+        return scale