autogluon.timeseries 1.4.1b20250906__py3-none-any.whl → 1.4.1b20251210__py3-none-any.whl

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

@@ -0,0 +1,305 @@
+# 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
+
+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: int | None = 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: int | None = 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

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

@@ -0,0 +1,333 @@
+# 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 cast
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+from .attention import (
+    AttentionAxis,
+    MultiHeadAttention,
+    SpaceWiseMultiheadAttention,
+    TimeWiseMultiheadAttention,
+)
+from .kvcache import KVCache
+from .rope import TimeAwareRotaryEmbedding
+from .rotary_embedding_torch import RotaryEmbedding
+
+
+class SwiGLU(torch.nn.Module):
+    """
+    https://arxiv.org/abs/2002.05202
+    NOTE: x should be 2x the size you want
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Note this ordering is unusual, but is done so to match xFormers
+        gate, x = x.chunk(2, dim=-1)
+        return F.silu(gate) * x
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, include_weight: bool = True, eps: float = 1e-8):
+        super(RMSNorm, self).__init__()
+        self.eps = eps
+        if include_weight:
+            self.scale: torch.nn.Parameter | None = torch.nn.Parameter(torch.ones(dim))
+        else:
+            self.scale = None
+
+    def forward(self, x: torch.Tensor):
+        x_normed = x / torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+        return x_normed if self.scale is None else x_normed * self.scale
+
+    def increment_and_forward_(self, x: torch.Tensor, y: torch.Tensor):
+        """
+        If you need the fused addition with RMS norm, do the same check here.
+        """
+        return self.forward(x + y)
+
+
+def make_batched_block_mask(t: torch.Tensor) -> torch.Tensor:
+    unsqueezed = rearrange(t, "... d -> ... 1 d")
+    return unsqueezed == unsqueezed.transpose(-1, -2)
+
+
+class TransformerLayer(torch.nn.Module):
+    """
+    A transformer block that applies multihead attention followed by a feedforward network.
+
+    The transformer can be configured to apply time-wise attention (i.e. attention over the time axis)
+    or space-wise attention (i.e. attention over the variate axis).
+
+    The transformer block uses pre-norm, which is a variant of the transformer architecture where
+    LayerNorm is applied before each sublayer, rather than after. This is the approach taken in
+    LLaMA and other recent transformer-based models.
+
+    The transformer block also uses SwiGLU, which is a variant of the Gated Linear Unit (GLU) activation
+    function. SwiGLU is a variant of the GLU activation that uses the Swish activation function. This
+    activation function has been used extensively in recent transformer-based models and has been shown
+    to improve performance.
+    """
+
+    embed_dim: int
+    num_heads: int
+    mlp_hidden_dim: int
+    dropout: float
+    attention_axis: AttentionAxis
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        mlp_hidden_dim: int,
+        dropout: float,
+        rotary_emb: RotaryEmbedding | None = None,
+        attention_axis: AttentionAxis = AttentionAxis.TIME,
+        RMS_norm: bool = True,
+        use_memory_efficient_attention: bool = True,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.mlp_hidden_dim = mlp_hidden_dim
+        self.dropout = dropout
+        self.attention_axis = attention_axis
+
+        if RMS_norm:
+            self.norm1: RMSNorm | torch.nn.LayerNorm = RMSNorm(embed_dim)
+            self.norm2: RMSNorm | torch.nn.LayerNorm = RMSNorm(embed_dim)
+
+        else:
+            self.norm1 = torch.nn.LayerNorm(embed_dim)
+            self.norm2 = torch.nn.LayerNorm(embed_dim)
+
+        self.attention: MultiHeadAttention
+
+        if attention_axis == AttentionAxis.TIME:
+            self.attention = TimeWiseMultiheadAttention(
+                embed_dim=embed_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                rotary_emb=rotary_emb,  # type: ignore
+                use_memory_efficient_attention=use_memory_efficient_attention,
+            )
+        elif attention_axis == AttentionAxis.SPACE:
+            self.attention = SpaceWiseMultiheadAttention(
+                embed_dim=embed_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                rotary_emb=None,
+                use_memory_efficient_attention=use_memory_efficient_attention,
+            )
+        else:
+            raise ValueError("Invalid attention axis")
+
+        self.mlp = torch.nn.Sequential(
+            torch.nn.Linear(embed_dim, 2 * mlp_hidden_dim),
+            SwiGLU(),
+            torch.nn.Linear(mlp_hidden_dim, embed_dim),
+            torch.nn.Dropout(dropout),
+        )
+
+    def forward(
+        self,
+        layer_idx: int,
+        inputs: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        kv_cache: KVCache | None = None,
+    ) -> torch.Tensor:
+        pre_norm_1 = self.norm1(inputs)
+        hidden_state = inputs + self.attention(layer_idx, pre_norm_1, attention_mask, kv_cache).contiguous()
+
+        pre_norm_2 = self.norm2(hidden_state)
+        return hidden_state + self.mlp(pre_norm_2)
+
+
+class Transformer(torch.nn.Module):
+    """
+    A stack of transformer layers. The transformer alternates between time-wise and space-wise attention
+    to learn both temporal and cross-variate dependencies in the data.
+
+    Based on the intuition that time-wise attention is more important overall than space-wise attention
+    (because an individual variate is more likely to be correlated with itself across time than with other variates),
+    the transformer can be configured to apply space-wise attention less frequently than time-wise attention.
+    This is controlled by the `spacewise_every_n_layers` parameter, which specifies how many time-wise transformer
+    layers to apply between every space-wise transformer layer.
+
+    Parameters
+    ----------
+    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.
+    use_memory_efficient_attention
+        Whether to use memory-efficient attention. If True, the model will use the memory-efficient from xFormers.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        embed_dim: int,
+        num_heads: int,
+        mlp_hidden_dim: int,
+        dropout: float,
+        spacewise_every_n_layers: int,
+        spacewise_first: bool,
+        use_memory_efficient_attention: bool = True,
+    ):
+        super().__init__()
+
+        assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads."
+
+        self.rotary_emb = TimeAwareRotaryEmbedding(
+            embed_dim // num_heads,
+            use_xpos=True,
+            cache_if_possible=True,
+            seq_before_head_dim=use_memory_efficient_attention,
+        )
+        attention_axes = self._get_layer_types(num_layers, spacewise_every_n_layers, spacewise_first)
+
+        self.use_memory_efficient_attention = use_memory_efficient_attention
+
+        self.layers = torch.nn.ModuleList(
+            [
+                TransformerLayer(
+                    embed_dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_hidden_dim=mlp_hidden_dim,
+                    dropout=dropout,
+                    rotary_emb=self.rotary_emb,
+                    attention_axis=attention_axes[i],
+                    use_memory_efficient_attention=self.use_memory_efficient_attention,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+    def _get_mask(
+        self,
+        num_heads: int,
+        dtype: torch.dtype,
+        id_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Unified method to create and process space-wise masks.
+
+        Args:
+            mask_type: Type of mask to create ('spacewise').
+            seq_len: Total sequence length.
+            num_heads: Number of attention heads.
+            device: Device where the mask should be created.
+            dtype: Desired dtype for the bias tensor.
+            id_mask: Mask for variates (used for spacewise masks).
+
+        Returns:
+            Processed attention mask tensor with the correct shape for the given mask type.
+        """
+
+        if id_mask is None:
+            raise ValueError("id_mask must be provided for spacewise masks.")
+
+        # Create spacewise mask
+        mask = make_batched_block_mask(id_mask.transpose(-1, -2))
+
+        if self.use_memory_efficient_attention:
+            mask = self._pad_to_multiple(mask)
+            mask = mask.float().masked_fill(~mask, float("-inf")).masked_fill(mask, 0.0).to(dtype)
+
+        # Rearrange for space-wise attention
+        mask = rearrange(mask, "batch seq_len variate1 variate2 -> (batch seq_len) 1 variate1 variate2")
+        # Stack along num_heads dimension
+        return mask.expand(-1, num_heads, -1, -1).contiguous()
+
+    def _pad_to_multiple(
+        self,
+        tensor: torch.Tensor,
+        multiple: int = 8,
+        causal: bool = False,  # New flag to indicate causal mask extension
+    ) -> torch.Tensor:
+        """
+        Pads the last two dimensions of a tensor to be divisible by `multiple`.
+        For causal masks, the padded area is filled with the continued lower-triangular pattern,
+        rather than with zeros.
+        """
+        pad_amount = (multiple - tensor.shape[-1] % multiple) % multiple
+        if pad_amount > 0:
+            new_size = tensor.shape[-1] + pad_amount
+            if causal:
+                # Create a full causal mask for the new size.
+                full_mask = torch.tril(torch.ones((new_size, new_size), dtype=tensor.dtype, device=tensor.device))
+                # Preserve any modifications from the original mask (e.g., condition tokens in top-left)
+                full_mask[: tensor.shape[-1], : tensor.shape[-1]] = tensor
+                tensor = full_mask
+            else:
+                tensor = F.pad(tensor, (0, pad_amount, 0, pad_amount))
+        return tensor
+
+    def _get_layer_types(
+        self,
+        num_layers: int,
+        spacewise_every_n_layers: int,
+        spacewise_first: bool,
+    ) -> list[AttentionAxis]:
+        if spacewise_every_n_layers == -1:
+            return [AttentionAxis.TIME] * num_layers
+        assert num_layers % spacewise_every_n_layers == 0
+
+        block = [AttentionAxis.TIME] * (spacewise_every_n_layers - 1)
+
+        if spacewise_first:
+            block = [AttentionAxis.SPACE] + block
+        else:
+            block = block + [AttentionAxis.SPACE]
+
+        layer_types = block * (num_layers // spacewise_every_n_layers)
+
+        return layer_types
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        id_mask: torch.Tensor,
+        kv_cache: KVCache | None = None,
+    ) -> torch.Tensor:
+        batch, _, seq_len, _ = inputs.shape
+        # Get the sequence length by looking up a timewise layer in the kv cache.
+        # Regardless of whether spacewise is first in the stack, the layer
+        # at index 1 is always a timewise layer.
+        seq_len = (kv_cache.seq_len(1) if kv_cache else 0) + seq_len
+
+        num_heads: int = cast(int, self.layers[0].num_heads)
+
+        timewise_attention_mask = None
+
+        # We create a space-wise ID mask by creating a block triangular mask from the ID mask
+        # in the space-wise direction. This ensures that the model can only attend to
+        # variates in the same group.
+        spacewise_attention_mask = self._get_mask(
+            num_heads=num_heads,
+            dtype=inputs.dtype,
+            id_mask=id_mask,
+        )
+
+        for layer_idx, layer in enumerate(self.layers):
+            inputs = layer(
+                layer_idx,
+                inputs,
+                (timewise_attention_mask if layer.attention_axis == AttentionAxis.TIME else spacewise_attention_mask),
+                kv_cache,
+            )
+        return inputs