Fix NaN bugs in the distributed calculation of mean and variance

[maliesen]

Due Diligence

• General:
  • 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

Issue/s resolved: There was a bug producing nan if the first two ranks are empty as described in #2495. Also, there is bug leading to nan if there is one rank with only one single element if one calls ht.var(..., axis=0, ddof=1).

Changes proposed:

• Make __merge_moments not produce any nans by not dividing by (n-ddof) and only divide by n if its safe (which is guaranteed by the early return if n2 == 0)
• Divide by (n-ddof) only once at the end of the calculation in the spirit of
  [1] J. Bennett, R. Grout, P. Pebay, D. Roe, D. Thompson, Numerically stable, single-pass, parallel statistics algorithms, IEEE International Conference on Cluster Computing and Workshops, 2009, Oct 2009, New Orleans, LA, USA.

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

[maliesen]

mpirun -n 1 pytest -vv passes locally so this might be a CI pipeline failure (exit code 15). Maybe we can just retry the job?

[brownbaerchen]

Looks really good, thanks!

The problem I have is that I am not an expert on this code and that the tests for var are a bit too confusing for me to just accept that surely everything works. That's not your fault, of course.
Before merging this, I would like to replace the tests that make sure the variance of ones is zero with tests that var on random data is the same between heat and numpy. In particular, with all this correction stuff which is very confusing.
If you have time to add that, that would be lovely. Otherwise, I can also do that in a separate PR and then we can merge this once all those tests pass on this branch.

[brownbaerchen]

Suggested change

	n = x.lnumel
	n = x.larray.numel

I would like to deprecate the lnumel function at some point because I think it's cleaner to operate on the local array directly when doing anything local.