diff --git a/heat/core/statistics.py b/heat/core/statistics.py index 3ce8dd4791..c2a9b43b66 100644 --- a/heat/core/statistics.py +++ b/heat/core/statistics.py @@ -1091,9 +1091,7 @@ def median( DNDarray.median.__doc__ = median.__doc__ -def __merge_moments( - m1: torch.Tensor, m2: torch.Tensor, correction: bool = True -) -> Tuple[torch.Tensor, ...]: +def __merge_moments(m1: Tuple, m2: Tuple, correction: bool = True) -> Tuple[torch.Tensor, ...]: """ Merge two statistical moments. If the length of ``m1`` and ``m2`` (must be equal) is ``==3`` then the second moment (variance) @@ -1120,6 +1118,8 @@ def __merge_moments( if len(m1) != len(m2): raise ValueError(f"m1 and m2 must be same length, currently {len(m1)} and {len(m2)}") n1, n2 = m1[-1], m2[-1] + if n2.sum() == 0: + return m1 mu1, mu2 = m1[-2], m2[-2] n = n1 + n2 delta = mu2 - mu1 diff --git a/tests/core/test_statistics.py b/tests/core/test_statistics.py index 51811288bd..27bb793e48 100644 --- a/tests/core/test_statistics.py +++ b/tests/core/test_statistics.py @@ -1,6 +1,7 @@ import numpy as np import torch import os +import unittest from itertools import combinations from scipy import stats as ss @@ -1599,3 +1600,12 @@ def test_var(self): # edge case from #2374 self.assertEqual(ht.var(ht.array([0.], split=None), axis=0, ddof=0), 0) + + @unittest.skipUnless(ht.communication.MPI_WORLD.size >= 3, "Test requires at least 3 tasks") + def test_first_two_leading_ranks_empty(self): + comm = self.comm + local_data = torch.tensor([], dtype=torch.float32) if comm.rank < 2 else torch.tensor([comm.rank], dtype=torch.float32) + data = ht.DNDarray(local_data, gshape=(comm.size-2,), dtype=ht.float32, split=0, device=ht.devices.cpu, comm=comm, balanced=False) + + self.assertEqual(ht.mean(data), np.mean(data.numpy())) + self.assertEqual(ht.var(data), np.var(data.numpy()))