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

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

@@ -0,0 +1,342 @@
+# Source: https://github.com/lucidrains/rotary-embedding-torch
+# MIT License
+#
+# Copyright (c) 2021 Phil Wang
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+from math import pi
+from typing import Literal
+
+import torch
+from einops import rearrange, repeat
+from torch import Tensor, broadcast_tensors, einsum, is_tensor, nn, tensor
+from torch.amp import autocast
+from torch.nn import Module
+
+# helper functions
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def slice_at_dim(t, dim_slice: slice, *, dim):
+    dim += t.ndim if dim < 0 else 0
+    colons = [slice(None)] * t.ndim
+    colons[dim] = dim_slice
+    return t[tuple(colons)]
+
+
+# rotary embedding helper functions
+
+
+def rotate_half(x):
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+@autocast("cuda", enabled=False)
+def apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2, freqs_seq_dim=None):
+    dtype = t.dtype
+
+    if not exists(freqs_seq_dim):
+        if freqs.ndim == 2 or t.ndim == 3:
+            freqs_seq_dim = 0
+
+    if t.ndim == 3 or exists(freqs_seq_dim):
+        seq_len = t.shape[seq_dim]
+        freqs = slice_at_dim(freqs, slice(-seq_len, None), dim=freqs_seq_dim)
+
+    rot_dim = freqs.shape[-1]
+    end_index = start_index + rot_dim
+
+    assert rot_dim <= t.shape[-1], (
+        f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}"
+    )
+
+    # Split t into three parts: left, middle (to be transformed), and right
+    t_left = t[..., :start_index]
+    t_middle = t[..., start_index:end_index]
+    t_right = t[..., end_index:]
+
+    # Apply rotary embeddings without modifying t in place
+    t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale)
+
+    out = torch.cat((t_left, t_transformed, t_right), dim=-1)
+
+    return out.type(dtype)
+
+
+# learned rotation helpers
+
+
+def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None):
+    if exists(freq_ranges):
+        rotations = einsum("..., f -> ... f", rotations, freq_ranges)
+        rotations = rearrange(rotations, "... r f -> ... (r f)")
+
+    rotations = repeat(rotations, "... n -> ... (n r)", r=2)
+    return apply_rotary_emb(rotations, t, start_index=start_index)
+
+
+# classes
+
+
+class RotaryEmbedding(Module):
+    def __init__(
+        self,
+        dim,
+        custom_freqs: Tensor | None = None,
+        freqs_for: Literal["lang", "pixel", "constant"] = "lang",
+        theta=10000,
+        max_freq=10,
+        num_freqs=1,
+        learned_freq=False,
+        use_xpos=False,
+        xpos_scale_base=512,
+        interpolate_factor=1.0,
+        theta_rescale_factor=1.0,
+        seq_before_head_dim=False,
+        cache_if_possible=True,
+        cache_max_seq_len=8192,
+    ):
+        super().__init__()
+        # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+        # has some connection to NTK literature
+        # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+
+        theta *= theta_rescale_factor ** (dim / (dim - 2))
+
+        self.freqs_for = freqs_for
+
+        if exists(custom_freqs):
+            freqs = custom_freqs
+        elif freqs_for == "lang":
+            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+        elif freqs_for == "pixel":
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == "constant":
+            freqs = torch.ones(num_freqs).float()
+
+        self.cache_if_possible = cache_if_possible
+        self.cache_max_seq_len = cache_max_seq_len
+
+        self.register_buffer("cached_freqs", torch.zeros(cache_max_seq_len, dim), persistent=False)
+        self.cached_freqs_seq_len = 0
+
+        self.freqs = nn.Parameter(freqs, requires_grad=learned_freq)
+
+        self.learned_freq = learned_freq
+
+        # dummy for device
+
+        self.register_buffer("dummy", torch.tensor(0), persistent=False)
+
+        # default sequence dimension
+
+        self.seq_before_head_dim = seq_before_head_dim
+        self.default_seq_dim = -3 if seq_before_head_dim else -2
+
+        # interpolation factors
+
+        assert interpolate_factor >= 1.0
+        self.interpolate_factor = interpolate_factor
+
+        # xpos
+
+        self.use_xpos = use_xpos
+
+        if not use_xpos:
+            return
+
+        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
+        self.scale_base = xpos_scale_base
+
+        self.register_buffer("scale", scale, persistent=False)
+        self.register_buffer("cached_scales", torch.zeros(cache_max_seq_len, dim), persistent=False)
+        self.cached_scales_seq_len = 0
+
+        # add apply_rotary_emb as static method
+
+        self.apply_rotary_emb = staticmethod(apply_rotary_emb)
+
+    @property
+    def device(self):
+        return self.dummy.device
+
+    def get_seq_pos(self, seq_len, device=None, dtype=None, offset=0):
+        device = default(device, self.device)
+        dtype = default(dtype, self.cached_freqs.dtype)
+
+        return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor
+
+    def rotate_queries_or_keys(self, t, seq_dim=None, offset=0, scale=None):
+        seq_dim = default(seq_dim, self.default_seq_dim)
+
+        assert not self.use_xpos or exists(scale), (
+            "you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings"
+        )
+
+        device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim]
+
+        seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset)
+
+        freqs = self.forward(seq, seq_len=seq_len, offset=offset)
+
+        if seq_dim == -3:
+            freqs = rearrange(freqs, "n d -> n 1 d")
+
+        return apply_rotary_emb(freqs, t, scale=default(scale, 1.0), seq_dim=seq_dim)
+
+    def rotate_queries_with_cached_keys(self, q, k, seq_dim=None, offset=0):
+        dtype, device, seq_dim = q.dtype, q.device, default(seq_dim, self.default_seq_dim)
+
+        q_len, k_len = q.shape[seq_dim], k.shape[seq_dim]
+        assert q_len <= k_len
+
+        q_scale = k_scale = 1.0
+
+        if self.use_xpos:
+            seq = self.get_seq_pos(k_len, dtype=dtype, device=device)
+
+            q_scale = self.get_scale(seq[-q_len:]).type(dtype)
+            k_scale = self.get_scale(seq).type(dtype)
+
+        rotated_q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset)
+        rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale**-1)
+
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+
+        return rotated_q, rotated_k
+
+    def rotate_queries_and_keys(self, q, k, seq_dim=None):
+        seq_dim = default(seq_dim, self.default_seq_dim)
+
+        assert self.use_xpos
+        device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim]
+
+        seq = self.get_seq_pos(seq_len, dtype=dtype, device=device)
+
+        freqs = self.forward(seq, seq_len=seq_len)
+        scale = self.get_scale(seq, seq_len=seq_len).to(dtype)
+
+        if seq_dim == -3:
+            freqs = rearrange(freqs, "n d -> n 1 d")
+            scale = rearrange(scale, "n d -> n 1 d")
+
+        rotated_q = apply_rotary_emb(freqs, q, scale=scale, seq_dim=seq_dim)
+        rotated_k = apply_rotary_emb(freqs, k, scale=scale**-1, seq_dim=seq_dim)
+
+        rotated_q = rotated_q.type(q.dtype)
+        rotated_k = rotated_k.type(k.dtype)
+
+        return rotated_q, rotated_k
+
+    def get_scale(self, t: Tensor, seq_len: int | None = None, offset=0):
+        assert self.use_xpos
+
+        should_cache = self.cache_if_possible and exists(seq_len) and (offset + seq_len) <= self.cache_max_seq_len
+
+        if should_cache and exists(self.cached_scales) and (seq_len + offset) <= self.cached_scales_seq_len:
+            return self.cached_scales[offset : (offset + seq_len)]
+
+        scale = 1.0
+        if self.use_xpos:
+            power = (t - len(t) // 2) / self.scale_base
+            scale = self.scale ** rearrange(power, "n -> n 1")
+            scale = repeat(scale, "n d -> n (d r)", r=2)
+
+        if should_cache and offset == 0:
+            self.cached_scales[:seq_len] = scale.detach()
+            self.cached_scales_seq_len = seq_len
+
+        return scale
+
+    def get_axial_freqs(self, *dims, offsets: (tuple[int | float, ...] | Tensor | None) = None):
+        Colon = slice(None)
+        all_freqs = []
+
+        # handle offset
+
+        if exists(offsets):
+            if not is_tensor(offsets):
+                offsets = tensor(offsets)
+
+            assert len(offsets) == len(dims)
+
+        # get frequencies for each axis
+
+        for ind, dim in enumerate(dims):
+            offset = 0
+            if exists(offsets):
+                offset = offsets[ind]
+
+            if self.freqs_for == "pixel":
+                pos = torch.linspace(-1, 1, steps=dim, device=self.device)
+            else:
+                pos = torch.arange(dim, device=self.device)
+
+            pos = pos + offset
+
+            freqs = self.forward(pos, seq_len=dim)
+
+            all_axis = [None] * len(dims)
+            all_axis[ind] = Colon
+
+            new_axis_slice = (Ellipsis, *all_axis, Colon)
+            all_freqs.append(freqs[new_axis_slice])
+
+        # concat all freqs
+
+        all_freqs = broadcast_tensors(*all_freqs)
+        return torch.cat(all_freqs, dim=-1)
+
+    @autocast("cuda", enabled=False)
+    def forward(self, t: Tensor, seq_len: int | None = None, offset=0):
+        should_cache = (
+            self.cache_if_possible
+            and not self.learned_freq
+            and exists(seq_len)
+            and self.freqs_for != "pixel"
+            and (offset + seq_len) <= self.cache_max_seq_len
+        )
+
+        if should_cache and exists(self.cached_freqs) and (offset + seq_len) <= self.cached_freqs_seq_len:
+            return self.cached_freqs[offset : (offset + seq_len)].detach()
+
+        freqs = self.freqs
+
+        freqs = einsum("..., f -> ... f", t.type(freqs.dtype), freqs)
+        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
+
+        if should_cache and offset == 0:
+            self.cached_freqs[:seq_len] = freqs.detach()
+            self.cached_freqs_seq_len = seq_len
+
+        return freqs

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