speechbrain
diff --git a/‎speechbrain/processing/features.py‎
Lines changed: 162 additions & 0 deletions b/‎speechbrain/processing/features.py‎
Lines changed: 162 additions & 0 deletions
@@ -32,11 +32,14 @@
 
 Authors
  * Mirco Ravanelli 2020
+ * Rogier van Dalen 2025
 """
 
 import math
+from typing import Tuple, Union
 
 import torch
+from torch.distributed import ReduceOp
 
 from speechbrain.dataio.dataio import length_to_mask
 from speechbrain.utils.checkpoints import (
@@ -993,6 +996,165 @@ def forward(self, x):
         return cw_x
 
 
+def gaussian_statistics(x: torch.Tensor, dim: Union[int, tuple, None] = None):
+    """
+    Compute first- and second-order moments of data, and return them as the
+    count, mean, and variance of a vector over one or more dimensions.
+
+    Arguments
+    ---------
+    x: torch.Tensor
+        The tensor to compute the statistics over
+    dim: int | tuple | None
+        The dimension or dimensions that the statistics should be computed over.
+        The other dimensions are retained in the output.
+        If None, then scalar-valued statistics will be returned.
+
+    Returns
+    -------
+    count
+        The number of sub-vectors or sub-tensors that the statistics were
+        computed over.
+    mean
+        The mean.
+    variance
+        The variance.
+    """
+
+    if dim is None:
+        number = math.prod(x.shape)
+    elif isinstance(dim, int):
+        number = x.shape[dim]
+    else:
+        assert isinstance(dim, tuple)
+        if dim == ():
+            return 1, x, torch.zeros_like(x)
+        number = 1
+        for d in dim:
+            number *= x.shape[d]
+
+    # First keep the dimensions so that broadcasting works.
+    mean_with_dims = torch.mean(x, dim=dim, keepdim=True)
+    mean = (
+        torch.squeeze(mean_with_dims)
+        if dim is None
+        else torch.squeeze(mean_with_dims, dim=dim)
+    )
+    variance = torch.mean(torch.square(x - mean_with_dims), dim=dim)
+
+    return (number, mean, variance)
+
+
+def combine_gaussian_statistics(
+    left_statistics: Tuple[int, torch.Tensor, torch.Tensor],
+    right_statistics: Tuple[int, torch.Tensor, torch.Tensor],
+):
+    """
+    Combine the first- and second-order moments from two pieces of data.
+    The data and the result is in the form (count, mean, variance).
+    The result is the mean and variance as if they have been computed on the
+    concatenation of the data for left_statistics and the data for
+    right_statistics.
+
+    Arguments
+    ---------
+    left_statistics: Tuple[int, torch.Tensor, torch.Tensor]
+        One set of statistics.
+    right_statistics: Tuple[int, torch.Tensor, torch.Tensor]
+        Another set of statistics.
+
+    Returns
+    -------
+    count
+        The total number of elements in the data.
+    mean
+        The combined mean.
+    variance
+        The combined variance, relative to the new mean.
+    """
+    left_count, left_mean, left_variance = left_statistics
+    right_count, right_mean, right_variance = right_statistics
+    assert left_mean.shape == left_variance.shape
+    assert left_mean.shape == right_mean.shape
+    assert left_variance.shape == right_variance.shape
+
+    count = left_count + right_count
+
+    left_weight = left_count / count
+    right_weight = right_count / count
+
+    mean = left_weight * left_mean + right_weight * right_mean
+
+    # Reconstruct the left and right variances relative to "mean".
+    compensated_left_variance = left_variance + torch.square(mean - left_mean)
+    compensated_right_variance = right_variance + torch.square(
+        mean - right_mean
+    )
+
+    variance = (
+        left_weight * compensated_left_variance
+        + right_weight * compensated_right_variance
+    )
+
+    return count, mean, variance
+
+
+def combine_gaussian_statistics_distributed(
+    statistics: Tuple[int, torch.Tensor, torch.Tensor],
+):
+    """
+    Combine the first- and second-order moments from multiple pieces of data
+    using torch.distributed.
+    The data and the result is in the form (count, mean, variance).
+    The result is the mean and variance as if they have been computed on the
+    concatenation of the data for statistics for all parallel processes.
+
+    Arguments
+    ---------
+    statistics: Tuple[int, torch.Tensor, torch.Tensor]
+        The new statistics for this process, to be combined with the current
+        statistics and the new statistics for all other processes.
+
+    Returns
+    -------
+    count
+        The total number of elements in the data across processes.
+    mean
+        The combined mean.
+    variance
+        The combined variance, relative to the new mean.
+    """
+    # This is the DDP version of combine_gaussian_statistics above.
+    local_count, local_mean, local_variance = statistics
+    global_count = ddp_all_reduce(torch.tensor(local_count), ReduceOp.SUM)
+
+    local_weight = local_count / global_count
+    global_mean = ddp_all_reduce(local_weight * local_mean, ReduceOp.SUM)
+
+    compensated_local_variance = local_variance + torch.square(
+        local_mean - global_mean
+    )
+    global_variance = ddp_all_reduce(
+        local_weight * compensated_local_variance, ReduceOp.SUM
+    )
+
+    return (global_count, global_mean, global_variance)
+
+
+def mean_std_update(x, mask, dim, run_count, run_mean, run_std=None):
+    assert torch.all(mask), "Not implemented yet"
+
+    # TODO implement run_std is None
+    current_statistics = (run_count, run_mean, torch.square(run_std))
+    new_statistics = combine_gaussian_statistics_distributed(
+        gaussian_statistics(x, dim=dim)
+    )
+    (count, mean, variance) = combine_gaussian_statistics(
+        current_statistics, new_statistics
+    )
+    return count, mean, torch.sqrt(variance)
+
+
 @register_checkpoint_hooks
 class InputNormalization(torch.nn.Module):
     """Performs mean and variance normalization of the input tensor.