Fix bug leading to mean and var being NaN if calculated distributed

[maliesen]

Due Diligence

• General:
  • [ X] title of the PR is suitable to appear in the Release Notes
• Implementation:
  • unit tests: all split configurations tested
  • unit tests: multiple dtypes tested
  • NEW unit tests: MPS tested (1 MPI process, 1 GPU)
  • benchmarks: created for new functionality
  • benchmarks: performance improved or maintained
  • documentation updated where needed

Description

If the data is distributed so that the first two leading ranks are empty a divison by zero in merge_moments leads to mean and variance being calculated as nan. This can be circumvented by early returning what has been merged so far if there is no new actual data coming in.

Issue/s resolved: #

Changes proposed:

Type of change

Bug fix (non-breaking change which fixes an issue)

Memory requirements

Performance

Does this change modify the behaviour of other functions? If so, which?

yes / no

[brownbaerchen]

Thanks a lot for spotting this! I have some minor comments but the overall fix is good.

[maliesen]

Thanks for the review!

[maliesen]

I put the pr in draft made as I realized the fix as is would introduce the following regression:

    @unittest.skipUnless(ht.communication.MPI_WORLD.size == 2, "Test for two tasks")
    def test_corrected_var_single_element(self):
        comm = self.comm
        ltensor_empty = torch.tensor([], dtype=torch.float32)
        ltensor_data = torch.tensor([0.], dtype=torch.float32)

        ldata_on_first_rank = ltensor_data if comm.rank == 0 else ltensor_empty
        data_on_first_rank = ht.DNDarray(ldata_on_first_rank, gshape=(2,), dtype=ht.float32, split=0, device=ht.devices.cpu, comm=comm, balanced=False)

        self.assertTrue(ht.isnan(ht.var(data_on_first_rank, ddof=1)))

This is because if the data (single element) is on the first rank, torch will give nan in https://github.com/helmholtz-analytics/heat/blob/main/heat/core/statistics.py#L2226 which will then be normalized to 0 in https://github.com/helmholtz-analytics/heat/blob/main/heat/core/statistics.py#L2228-L2229 so the merge_moments function is responsible for turning this in nan at https://github.com/helmholtz-analytics/heat/blob/main/heat/core/statistics.py#L1132 which would not happen anymore by the early return in this fix.

The distributed calculation of the corrected variance has a nan bug already before this fix (in the __moment_w_axis code branch) and I planned on fixing this as a separate pr but due to the regression I think it is better to do it at once. Therefore I merged the fix of this pr into #2507 and if that pr is successful this pr may be closed.

[brownbaerchen]

Closing as superseded by #2507.