Vectorize and unify running_stats computation

sdm/models/kumo/timeseries/forecast/model.py

@@ -328,7 +328,8 @@ def forward(
         return self.decode(encoded, mean, scale)
 
     def decode(self, encoded: Tensor, mean: Tensor, scale: Tensor) -> Tensor:
-        return self.revin.inverse(self.head(encoded.flatten(-2)), mean, scale)
+        out = self.head(encoded.flatten(-2)).float()
+        return self.revin.inverse(out, mean, scale)
 
     def encode_memory(
         self,
@@ -362,7 +363,7 @@ def encode(
         self,
         x: Tensor,  # [..., C, R]
     ) -> tuple[Tensor, Tensor, Tensor]:
-        mean, scale = self.revin.compute_stats(x)
+        mean, scale = self.revin.stats(x, correction=0)
         x = self.revin(x, mean, scale)
         x = x.nan_to_num(nan=0, posinf=0, neginf=0)
         x = x.unflatten(-1, (-1, self.patch_size))  # [..., C, N, P]

sdm/models/timesfm3/model.py

@@ -37,8 +37,6 @@
 from sdm.models.timesfm3.util import (
     crossfade_patches,
     gather_future_patches,
-    get_running_stats,
-    update_running_stats,
 )
 from sdm.nn import RevIN
 from sdm.tensor.table import TableSchema
@@ -280,17 +278,17 @@ def forward(
         # Gather future patches and normalize:
         mask = x.isnan()
         x[mask] = 0.0
-        count, mean, std = get_running_stats(x, mask)
+        mean, std, count = self.revin.running_stats(x, mask, correction=0)
 
         future_x, past_end = gather_future_patches(x, num_patches=O // P)
         future_mask, _ = gather_future_patches(mask, num_patches=O // P)
         future_mask |= past_end
         # Exclude future information for context only features and targets:
         future_mask[..., C:, :, :] = True
 
-        scale = std.where(std >= 1e-6, 1.0).unsqueeze(-1)
-        x = self.revin(x, mean.unsqueeze(-1), scale)
-        future_x = self.revin(future_x, mean.unsqueeze(-1), scale)
+        scale = std.where(std >= 1e-6, 1.0)
+        x = self.revin(x, mean, scale)
+        future_x = self.revin(future_x, mean, scale)
         x[mask] = 0.0
         future_x[future_mask] = 0.0
 
@@ -307,48 +305,43 @@ def forward(
 
         x = self.patch_embedding(x)  # [..., C, N, D]
         out = self.icl_block(x, patch_mask=patch_mask)  # [..., C, N, O * Q]
-        out = out.unflatten(-1, (self.output_patch_size, -1))
+        out = out.unflatten(-1, (self.output_patch_size, -1)).float()
         first = mid // P - 1
 
         # Refine forecast normalization statistics using predicted medians:
         if self.refinement and horizon > 0:
             median = out[..., out.size(-1) // 2].unflatten(-1, (O // P, P))
-            mask = torch.zeros_like(median[..., 0, 0, :], dtype=torch.bool)
-
-            # Get most recent context statistics:
-            carry_count = count[..., first]
-            carry_mean = mean[..., first]
-            carry_std = std[..., first]
+            state = (  # Get most recent context statistics:
+                mean[..., first : first + 1, :],
+                std[..., first : first + 1, :],
+                count[..., first : first + 1, :],
+            )
 
             means, stds = [], []
             for i in range(first, out.size(-3) - 1, O // P):
                 anchor = self.revin.inverse(  # Denormalize output patch:
                     median[..., i, :, :],
-                    carry_mean[..., None, None],
-                    carry_std[..., None, None],
+                    mean=state[0],
+                    scale=state[1],
                 ).clamp(-1e20, 1e20)  # [..., C, O / P, P]
 
                 for j in range(min(O // P, out.size(-3) - i - 1)):
                     # Update running statistics:
-                    carry_count, carry_mean, carry_std = update_running_stats(
-                        count=carry_count,
-                        mean=carry_mean,
-                        std=carry_std,
-                        x=anchor[..., j, :],
-                        mask=mask,
+                    state = self.revin.running_stats(
+                        x=anchor[..., j : j + 1, :],
+                        state=state,
+                        correction=0,
                     )
-                    means.append(carry_mean)
-                    stds.append(carry_std)
+                    means.append(state[0])
+                    stds.append(state[1])
 
-            means = [mean[..., : first + 1], torch.stack(means, dim=-1)]
-            mean = torch.cat(means, dim=-1)
-            stds = [std[..., : first + 1], torch.stack(stds, dim=-1)]
-            std = torch.cat(stds, dim=-1)
+            mean = torch.cat([mean[..., : first + 1, :], *means], dim=-2)
+            std = torch.cat([std[..., : first + 1, :], *stds], dim=-2)
 
         out = self.revin.inverse(
             out,
-            mean[..., None, None],
-            scale=std[..., None, None],
+            mean=mean.unsqueeze(-1),
+            scale=std.unsqueeze(-1),
         ).clamp(-1e20, 1e20)
 
         # Narrow target output patches:

sdm/models/timesfm3/util.py

@@ -19,85 +19,6 @@
 from torch import Tensor
 
 
-def update_running_stats(
-    count: Tensor,  # [..., C]
-    mean: Tensor,  # [..., C]
-    std: Tensor,  # [..., C]
-    x: Tensor,  # [..., C, P]
-    mask: Tensor,  # [..., C, P]
-) -> tuple[Tensor, Tensor, Tensor]:  # [..., C]
-    """Update running statistics for a new patch.
-
-    Args:
-        count: Count of valid values with shape ``[..., C]``, where ``C`` is
-            the number of columns.
-        mean: Running mean with shape ``[..., C]``.
-        std: Running standard deviation with shape ``[..., C]``.
-        x: New input with shape ``[..., C, P]``, where ``P`` is the patch size.
-        mask: Invalid-value mask with shape ``[..., C, P]``.
-
-    Returns:
-        Updated count, mean, and standard deviation.
-    """
-    valid = ~mask
-
-    inc_count = valid.sum(dim=-1)
-    inc_sum = torch.where(valid, x, 0.0).sum(dim=-1)
-    inc_mean = inc_sum / inc_count.clamp(min=1)
-    tmp = (x - inc_mean.unsqueeze(-1)).square()
-    inc_var = torch.where(valid, tmp, 0.0).sum(dim=-1) / inc_count.clamp(min=1)
-    inc_std = inc_var.sqrt()
-
-    out_count = count + inc_count
-    out_mean = (count * mean + inc_count * inc_mean) / out_count.clamp(min=1)
-    out_var = (
-        count * std * std
-        + inc_count * inc_std * inc_std
-        + count * (mean - out_mean) * (mean - out_mean)
-        + inc_count * (inc_mean - out_mean) * (inc_mean - out_mean)
-    ) / out_count.clamp(min=1)
-    out_std = out_var.sqrt()
-
-    return out_count, out_mean, out_std
-
-
-def get_running_stats(
-    x: Tensor,  # [..., C, N, P]
-    mask: Tensor,  # [..., C, N, P]
-) -> tuple[Tensor, Tensor, Tensor]:  # [..., C, N]
-    """Compute cumulative statistics patch by patch.
-
-    Args:
-        x: Input with shape ``[..., C, N, P]``, where ``C`` is the number of
-            columns, ``N`` is the number of patches, and ``P`` is the patch
-            size.
-        mask: Invalid-value mask with shape ``[..., C, N, P]``.
-
-    Returns:
-        Cumulative count, mean, and standard deviation.
-    """
-    *B, C, N, _ = x.size()
-
-    count = torch.zeros((*B, C), dtype=torch.int64, device=x.device)
-    mean = torch.zeros((*B, C), dtype=torch.float32, device=x.device)
-    std = torch.zeros((*B, C), dtype=torch.float32, device=x.device)
-
-    counts, means, stds = [], [], []
-    for i in range(N):
-        count, mean, std = update_running_stats(
-            count, mean, std, x[..., i, :], mask[..., i, :]
-        )
-        counts.append(count)
-        means.append(mean)
-        stds.append(std)
-
-    return (
-        torch.stack(counts, dim=-1),
-        torch.stack(means, dim=-1),
-        torch.stack(stds, dim=-1),
-    )
-
-
 def gather_future_patches(
     x: Tensor,  # [..., N, P]
     num_patches: int,

sdm/nn/revin.py

@@ -41,40 +41,112 @@ def __init__(
                 torch.zeros(affine_channels, 1, **factory_kwargs)
             )
 
-    def compute_stats(
-        self,
-        x: Tensor,  # [..., C, R]
-        mask: Tensor | None = None,  # [..., C, R]
-    ) -> tuple[Tensor, Tensor]:  # [..., C, 1]
+    @staticmethod
+    @torch.no_grad()
+    def stats(
+        x: Tensor,  # [..., R, D]
+        mask: Tensor | None = None,  # [..., R, D]
+        correction: int = 1,
+    ) -> tuple[Tensor, Tensor]:  # [..., R, 1]
         """Compute mean and standard deviation over the final dimension.
 
-        NaNs and masked values are excluded.
+        NaNs and masked values are excluded. Mean and standard deviation
+        are computed and returned in
+        :external+torch:ref:`torch.float64 <dtype-doc>`.
 
         Args:
-            x: Tensor with shape ``[..., C, R]``.
-            mask: Boolean observed-value mask with shape ``[..., C, R]``.
+            x: Tensor with shape ``[..., R, D]``.
+            mask: Boolean mask with shape ``[..., R, D]``. Values marked
+                ``True`` are excluded.
+            correction: Difference between the sample size and sample degrees
+                of freedom.
 
         Returns:
-            Mean and standard deviation with shape ``[..., C, 1]``.
+            Mean and standard deviation with shape ``[..., R, 1]``.
         """
-        x = x.detach()
         if mask is not None:
-            x = x.masked_fill(~mask, torch.nan)
+            x = x.masked_fill(mask, torch.nan)
 
+        x = x.double()
         mean = x.nanmean(dim=-1, keepdim=True)
-        std = (x - mean).square().nanmean(dim=-1, keepdim=True).sqrt()
+        var = (x - mean).square().nanmean(dim=-1, keepdim=True)
+        if correction != 0:
+            count = (~x.isnan()).sum(dim=-1, keepdim=True)
+            var = var * (count.double() / (count - correction).clamp(min=1))
+        std = var.sqrt()
 
         mask = mean.isnan()
         return mean.masked_fill(mask, 0.0), std.masked_fill(mask, 1.0)
 
+    @staticmethod
+    @torch.no_grad()
+    def running_stats(
+        x: Tensor,  # [..., R, D]
+        mask: Tensor | None = None,  # [..., R, D]
+        state: tuple[Tensor, Tensor, Tensor] | None = None,
+        correction: int = 1,
+    ) -> tuple[Tensor, Tensor, Tensor]:  # [..., R, 1]
+        """Compute cumulative statistics over patched sequences.
+
+        NaNs and masked values are excluded. Mean and standard deviation
+        are computed and returned in
+        :external+torch:ref:`torch.float64 <dtype-doc>`.
+
+        Args:
+            x: Tensor with shape ``[..., R, D]``.
+            mask: Boolean mask with shape ``[..., R, D]``. Values marked
+                ``True`` are excluded.
+            state: Previous mean, standard deviation, and count for streaming
+                statistics.
+            correction: Degrees of freedom subtracted from the observed count.
+
+        Returns:
+            Cumulative mean, standard deviation, and observed count for the
+            supplied sequence, each with shape ``[..., R, 1]``.
+        """
+        D = x.size(-1)
+        mask = x.isnan() if mask is None else x.isnan() | mask
+        mask = mask.flatten(-2)  # [..., R * D]
+        x = x.flatten(-2).double().masked_fill(mask, 0)  # [..., R * D]
+        count = (~mask).cumsum(dim=-1)
+        mean = x.cumsum(dim=-1) / count.clamp(min=1)
+
+        # Welford increments accumulate squared deviations over time.
+        prev = torch.nn.functional.pad(mean[..., :-1], (1, 0))
+        inc = ((x - prev) * (x - mean)).masked_fill(mask, 0)
+        m2 = inc.cumsum(dim=-1)
+
+        # Keep statistics at the final value of each patch.
+        count = count[..., D - 1 :: D].unsqueeze(-1).clone()
+        mean = mean[..., D - 1 :: D].unsqueeze(-1).clone()
+        m2 = m2[..., D - 1 :: D].unsqueeze(-1)
+
+        if state is not None:
+            state_mean, state_std, state_count = state
+            state_mean = state_mean[..., -1:, :]
+            state_std = state_std[..., -1:, :]
+            state_count = state_count[..., -1:, :]
+
+            total_count = state_count + count
+            weight = count.double() / total_count.clamp(min=1)
+            delta = mean - state_mean
+            m2 += state_std.square() * (state_count - correction).clamp(min=1)
+            m2 += delta.square() * state_count * weight
+            mean = state_mean + delta * weight
+            count = total_count
+
+        std = (m2 / (count - correction).clamp(min=1)).clamp(min=0).sqrt()
+        return mean, std, count
+
     def forward(self, x: Tensor, mean: Tensor, scale: Tensor) -> Tensor:
         """The forward pass.
 
         Args:
-            x: Tensor with shape ``[..., C, R]``.
-            mean: The mean with shape ``[..., C, 1]``.
-            scale: The scale with shape ``[..., C, 1]``.
+            x: Tensor with shape ``[..., R, D]``.
+            mean: The mean with shape ``[..., R, 1]``.
+            scale: The scale with shape ``[..., R, 1]``.
         """
+        mean, scale = mean.to(x.dtype), scale.to(x.dtype)
         out = (x - mean) / (scale + self.eps)
         if self.weight is not None and self.bias is not None:
             out = out * self.weight + self.bias
@@ -84,11 +156,11 @@ def inverse(self, x: Tensor, mean: Tensor, scale: Tensor) -> Tensor:
         """Invert the forward pass.
 
         Args:
-            x: Normalized values or forecasts. With affine parameters, shape
-            x: Tensor with shape ``[..., C, R]``.
-            mean: The mean with shape ``[..., C, 1]``.
-            scale: The scale with shape ``[..., C, 1]``.
+            x: Tensor with shape ``[..., R, D]``.
+            mean: The mean with shape ``[..., R, 1]``.
+            scale: The scale with shape ``[..., R, 1]``.
         """
+        mean, scale = mean.to(x.dtype), scale.to(x.dtype)
         if self.weight is not None and self.bias is not None:
             x = (x - self.bias) / (self.weight + self.eps * self.eps)
         return x * (scale + self.eps) + mean