autogluon.timeseries 1.4.1b20250902__py3-none-any.whl → 1.4.1b20251003__py3-none-any.whl

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,94 @@
+# 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 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
+
+
+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: Optional[int] = None,
+        seq_pos: Optional[torch.Tensor] = 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: Optional[int] = 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

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

@@ -0,0 +1,306 @@
+# 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 warnings
+from typing import Optional
+
+import torch
+from einops import repeat
+from gluonts.core.component import validated
+from gluonts.torch.scaler import Scaler
+
+
+def compute_causal_statistics(
+    data: torch.Tensor,
+    weights: torch.Tensor,
+    padding_mask: torch.Tensor,
+    dim: int,
+    minimum_scale: float,
+    use_bessel_correction: bool = True,
+    stabilize_with_global: bool = False,
+    scale_factor_exponent: float = 10.0,
+    prefix_length: Optional[int] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute causal mean and scale statistics along a specified dimension using
+    a vectorized implementation of Welford's algorithm for numerical stability.
+
+    This implementation avoids explicit loops while maintaining the numerical stability
+    of Welford's algorithm, achieving better performance with the same robustness
+    against overflow issues.
+
+
+    Can optionally use global statistics to stabilize causal statistics by clamping
+    extreme values, preventing instability while preserving a relaxed version of the
+    causal property. This allows a controlled amount of future information leakage,
+    introducing an explicit tradeoff between causality and stability.
+    extreme values, preventing instability while preserving the causal property.
+
+    Parameters
+    ----------
+    data
+        The input data tensor
+    weights
+        The weight tensor (same shape as data)
+    padding_mask
+        The padding mask tensor (same shape as data)
+    dim
+        The dimension along which to compute statistics (must be -1, the time dimension)
+    minimum_scale
+        Minimum scale value to use
+    use_bessel_correction
+        Whether to use Bessel's correction to get an unbiased estimator
+    stabilize_with_global
+        Whether to use global statistics to stabilize the causal statistics by clamping
+        extreme values
+    scale_factor_exponent
+        Exponent that controls the allowed range of deviation from global scale.
+        For example, with exponent=1.0, causal scale must be between 0.1x and 10x the global scale.
+        With exponent=2.0, the range would be 0.01x to 100x.
+    prefix_length
+        If specified, the global statistics will be computed using only the prefix length
+        requested. This is used for multistep decoding, where we only want to use the
+        initial historical data to compute the global statistics. If stabilize_with_global
+        is False, this parameter is ignored.
+
+    Returns
+    -------
+    tuple[torch.Tensor, torch.Tensor]
+        Causal mean and scale tensors, potentially stabilized with global statistics
+    """
+    # Assert that dim is -1 (last dimension)
+    assert dim == -1, "compute_causal_statistics only supports dim=-1 (last dimension)"
+
+    with torch.no_grad():
+        # Apply padding mask to weights
+        weights = weights * padding_mask
+
+        # Try to use higher precision for numerical stability
+        try:
+            high_precision_data = data.to(torch.float64)
+            high_precision_weights = weights.to(torch.float64)
+        except TypeError:
+            # Fallback for devices that don't support float64
+            warnings.warn(
+                f"Float64 is not supported by device {data.device}. "
+                "Using float32 instead for causal scaler calculations. "
+                "This may lead to numerical issues if the data contains extreme values.",
+                RuntimeWarning,
+            )
+            high_precision_data = data.to(torch.float32)
+            high_precision_weights = weights.to(torch.float32)
+
+        # Check if deterministic algorithms are enabled and we're using CUDA.
+        # Cumsum operations do not support deterministic mode in CUDA,
+        # so we need to disable it for just this section.
+        prev_deterministic = torch.are_deterministic_algorithms_enabled()
+        if prev_deterministic and data.device.type == "cuda":
+            # Disable deterministic algorithms for operations
+            torch.use_deterministic_algorithms(False)
+
+        try:
+            # Create weighted data
+            weighted_data = high_precision_weights * high_precision_data
+
+            # Compute cumulative sum of weights and weighted data along time dimension
+            cum_weights = torch.cumsum(high_precision_weights, dim=dim)
+            cum_values = torch.cumsum(weighted_data, dim=dim)
+
+            # Avoid division by zero for the first time step or when no valid values
+            denominator = cum_weights.clamp_min(1.0)
+
+            # Compute causal means at each time step
+            causal_means = cum_values / denominator
+
+            # For Welford's algorithm, we need to compute the correction term
+            # using the difference between the current value and the current mean
+
+            # Create shifted version of causal means to compute delta efficiently
+            # First item in shifted_means will be zero
+            shifted_means = torch.zeros_like(causal_means)
+            shifted_means[..., 1:] = causal_means[..., :-1]
+
+            # Compute delta between current data point and previous mean
+            # For t=0, this is just the first data point
+            delta = high_precision_data - shifted_means
+
+            # Compute the increment term for Welford's algorithm.
+            # This is defined as the product of the delta and the difference between the current data point and the causal mean.
+            # This is where we avoid the traditional E[X²] - E[X]² computation
+            increment = delta * (high_precision_data - causal_means) * high_precision_weights
+
+            # The Welford algorithm uses the term m_2, which is the cumulative sum of the increment term.
+            # This is an accumulator that helps us compute the second moment (hence m_2) of the distribution.
+            # Compute cumulative sum of the increment term
+            m_2 = torch.cumsum(increment, dim=dim)
+
+            # Compute variance according to Welford's algorithm
+            if use_bessel_correction:
+                causal_variance = m_2 / torch.clamp(denominator - 1.0, min=1.0)
+            else:
+                causal_variance = m_2 / denominator
+
+            # Add minimum scale but keep in high precision for now
+            causal_scale = torch.sqrt(causal_variance + minimum_scale)
+
+            # Apply stabilization with global statistics if requested
+            if stabilize_with_global:
+                if prefix_length is not None:
+                    # Create a prefix mask for global statistics computation
+                    prefix_mask = torch.zeros_like(weights)
+                    prefix_mask[..., :prefix_length] = 1.0
+
+                    # Apply prefix mask to restrict computation to prefix
+                    weighted_data = weighted_data * prefix_mask
+                    weights = weights * prefix_mask
+                    padding_mask = padding_mask * prefix_mask
+
+                # Calculate scale factors from the exponent
+                scale_factor_min = 10.0 ** (-scale_factor_exponent)
+                scale_factor_max = 10.0**scale_factor_exponent
+
+                global_denominator = (weights * padding_mask).sum(dim, keepdim=True).clamp_min(1.0)
+                global_means = (weighted_data).sum(dim, keepdim=True) / global_denominator
+                global_means = torch.nan_to_num(global_means)
+
+                global_variance = (((high_precision_data - global_means) * weights * padding_mask) ** 2).sum(
+                    dim, keepdim=True
+                ) / global_denominator
+                global_scale = torch.sqrt(global_variance + minimum_scale)
+
+                # Expand global statistics to match the time dimension
+                expanded_global_scale = global_scale.expand_as(causal_scale)
+
+                # Define bounds using scale factors
+                min_allowed_scale = expanded_global_scale * scale_factor_min
+                max_allowed_scale = expanded_global_scale * scale_factor_max
+
+                # Clamp the causal scale between min_allowed_scale and max_allowed_scale
+                causal_scale = torch.clamp(
+                    causal_scale,
+                    min=torch.max(torch.tensor(minimum_scale, device=causal_scale.device), min_allowed_scale),
+                    max=max_allowed_scale,
+                )
+
+            # Now convert means and scale to original dtype after all numerical operations
+            causal_means = causal_means.to(data.dtype)
+            causal_scale = causal_scale.to(data.dtype)
+
+        finally:
+            # Restore original deterministic setting if it was changed
+            if prev_deterministic and data.device.type == "cuda":
+                torch.use_deterministic_algorithms(True)
+
+        return causal_means, causal_scale
+
+
+class CausalPatchStdMeanScaler(Scaler):
+    """
+    Causally scales data in patches, where each patch uses statistics computed
+    from all data up to and including that patch. Within each patch, all timesteps
+    use the same scaling values.
+
+    This approach provides more stability than per-timestep causal scaling while
+    still maintaining the causal property (not using future data).
+
+    Can optionally stabilize causal statistics using global statistics to prevent
+    extreme values, while preserving the causal property.
+
+    The statistics are computed using Welford's algorithm, which provides better
+    numerical stability compared to the direct computation of variance, especially
+    when dealing with large values or a large number of data points.
+
+    Note: This scaler only works with the following constraints:
+    - The input must have shape [batch, variates, time_steps]
+    - It only operates on the last dimension (-1)
+    - The time_steps must be divisible by patch_size
+
+    Parameters
+    ----------
+    dim
+        dimension along which to compute the causal scale. Must be -1 (the last dimension).
+    patch_size
+        number of timesteps in each patch
+    minimum_scale
+        default scale that is used for elements that are constantly zero
+        along dimension `dim` or for the first patch.
+    use_bessel_correction
+        whether to use Bessel's correction to get an unbiased estimator
+    stabilize_with_global
+        whether to use global statistics to stabilize extreme causal statistics
+    scale_factor_exponent
+        exponent that controls the allowed range of deviation from global scale.
+        For example, with exponent=1.0, causal scale must be between 0.1x and 10x the global scale.
+        With exponent=2.0, the range would be 0.01x to 100x.
+    """
+
+    @validated()
+    def __init__(
+        self,
+        dim: int = -1,
+        patch_size: int = 32,
+        minimum_scale: float = 0.1,
+        use_bessel_correction: bool = True,
+        stabilize_with_global: bool = False,
+        scale_factor_exponent: float = 10.0,
+    ) -> None:
+        super().__init__()
+        assert dim == -1, "CausalPatchStdMeanScaler only supports dim=-1 (last dimension)"
+        self.dim = dim
+        self.patch_size = patch_size
+        self.minimum_scale = minimum_scale
+        self.use_bessel_correction = use_bessel_correction
+        self.stabilize_with_global = stabilize_with_global
+        self.scale_factor_exponent = scale_factor_exponent
+
+    def __call__(  # type: ignore[override]
+        self,
+        data: torch.Tensor,
+        padding_mask: torch.Tensor,
+        weights: torch.Tensor,
+        prefix_length: Optional[int] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        assert data.shape == weights.shape, "data and weights must have same shape"
+        assert len(data.shape) == 3, "Input data must have shape [batch, variates, time_steps]"
+
+        with torch.no_grad():
+            # Get the number of time steps (last dimension)
+            time_steps = data.shape[-1]
+
+            # Assert that time_steps is divisible by patch_size
+            assert time_steps % self.patch_size == 0, (
+                f"Time steps ({time_steps}) must be divisible by patch size ({self.patch_size})"
+            )
+
+            # First compute causal statistics with optional stabilization
+            causal_means, causal_scale = compute_causal_statistics(
+                data,
+                weights,
+                padding_mask,
+                -1,
+                self.minimum_scale,
+                self.use_bessel_correction,
+                self.stabilize_with_global,
+                self.scale_factor_exponent,
+                prefix_length,
+            )
+
+            # Unfold the causal means and scales to get the patches
+            means_unfolded = causal_means.unfold(-1, self.patch_size, self.patch_size)
+            scales_unfolded = causal_scale.unfold(-1, self.patch_size, self.patch_size)
+
+            # Get the last element of each patch (the most recent statistic)
+            patch_stats_means = means_unfolded[..., -1]
+            patch_stats_scales = scales_unfolded[..., -1]
+
+            # Tile the patch statistics across time dimension using einops.repeat
+            # With our fixed [batch, variates, num_patches] shape this is much simpler
+            patch_means = repeat(patch_stats_means, "b v p -> b v (p s)", s=self.patch_size)
+            patch_scales = repeat(patch_stats_scales, "b v p -> b v (p s)", s=self.patch_size)
+
+            # Apply normalization
+            scaled_data = (data - patch_means) / patch_scales
+
+            return scaled_data, patch_means, patch_scales