Merge pull request #341 from kashif/timesfmv1-var

[TimesFMv1] fix variance calculation

v1/src/timesfm/patched_decoder.py

@@ -190,13 +190,13 @@ def _masked_mean_std(inputs: JTensor,
 
   num_valid_elements = jnp.where(num_valid_elements == 0, 1, num_valid_elements)
 
-  # Calculate the masked sum and squared sum of M
+  # Calculate the masked sum for mean and centered squared sum for variance.
   masked_sum = jnp.sum(arr * mask, axis=1)
-  masked_squared_sum = jnp.sum((arr * mask)**2, axis=1)
 
   # Calculate the masked mean and standard deviation
   masked_mean = masked_sum / num_valid_elements
-  masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
+  centered = (arr - masked_mean[:, None]) * mask
+  masked_var = jnp.sum(centered**2, axis=1) / num_valid_elements
   masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var)
   masked_std = jnp.sqrt(masked_var)
 
@@ -295,7 +295,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
       patched_pads: JTensor) -> Tuple[JTensor, Tuple[JTensor, JTensor]]:
     """Input is of shape [B, N, P]."""
     mu, sigma = _masked_mean_std(inputs, patched_pads)
-    sigma = jnp.where(sigma < _TOLERANCE, 1.0, sigma)
+    sigma = jnp.maximum(sigma, _TOLERANCE)
     # Normalize each patch.
     outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
     outputs = jnp.where(

v1/src/timesfm/pytorch_patched_decoder.py

@@ -94,26 +94,16 @@ def _masked_mean_std(
 
   # Calculate the number of valid elements
   num_valid_elements = torch.sum(mask, dim=1)
-  num_valid_elements = torch.where(
-      num_valid_elements == 0,
-      torch.tensor(1,
-                   dtype=num_valid_elements.dtype,
-                   device=num_valid_elements.device),
-      num_valid_elements,
-  )
+  num_valid_elements = torch.clamp(num_valid_elements, min=1.0)
 
-  # Calculate the masked sum and squared sum
+  # Calculate the masked sum and mean
   masked_sum = torch.sum(arr * mask, dim=1)
-  masked_squared_sum = torch.sum((arr * mask)**2, dim=1)
-
-  # Calculate the masked mean and standard deviation
   masked_mean = masked_sum / num_valid_elements
-  masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
-  masked_var = torch.where(
-      masked_var < 0.0,
-      torch.tensor(0.0, dtype=masked_var.dtype, device=masked_var.device),
-      masked_var,
-  )
+
+  # Calculate the masked variance using centered values (numerically stable)
+  masked_centered_arr = (arr - masked_mean.unsqueeze(-1)) * mask
+  masked_var = torch.sum(masked_centered_arr**2, dim=1) / num_valid_elements
+  masked_var = torch.clamp(masked_var, min=0.0)
   masked_std = torch.sqrt(masked_var)
 
   return masked_mean, masked_std
@@ -616,11 +606,7 @@ class PatchedTimeSeriesDecoder(nn.Module):
   ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
     """Input is of shape [B, N, P]."""
     mu, sigma = _masked_mean_std(inputs, patched_pads)
-    sigma = torch.where(
-        sigma < self.config.tolerance,
-        torch.tensor(1.0, dtype=sigma.dtype, device=sigma.device),
-        sigma,
-    )
+    sigma = torch.clamp(sigma, min=self.config.tolerance)
 
     # Normalize each patch
     outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]