fix: use float64 accumulators in PearsonCorrelation to prevent catastrophic cancellation

[tejas-ae]

Fixes #3662.

The naive E[X²] - (E[X])² formula for variance is mathematically correct but numerically unstable. When values have large magnitude relative to their variance — for example an offset of 1e8 with small inter-sample differences — both E[X²] and (E[X])² are around 1e16 while their difference (the actual variance) is O(1). In float32, the unit in the last place at that scale is roughly 10^9, so the variance is completely lost and the metric returns 0.0.

The fix is to accumulate in float64 on devices that support it. MPS does not support float64, so it falls back to float32 and retains the previous behaviour. The public return type (Python float) is unchanged.

# before (float32 accumulators)
offset = 1e8
y_pred = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float64) + offset
y      = torch.tensor([2.0, 4.0, 6.0, 8.0, 10.0], dtype=torch.float64) + offset
m = PearsonCorrelation()
m.update((y_pred, y))
m.compute()  # 0.0  (wrong; expected 1.0)

# after (float64 accumulators)
m.compute()  # 1.0  ✓

A new test test_numerical_stability_large_offset covers this case. All existing non-distributed tests pass.

you know the codebase far better than I do — happy to adjust if a different approach (e.g., a shared Welford-based utility across PearsonCorrelation, R2Score, and FID as discussed in the issue) is preferred.

[siddhant-shahhh]

The bug is real and this fix works. One thing worth noting: the part that fixes it is the input cast, not the float64 accumulators.

I reproduced the accumulator math with float32 inputs at offset 1e4:

master (float32 accumulators)            r = 800000000.0
float64 accumulators, inputs NOT cast    r = 440000000.59604645
float64 accumulators + input cast (PR)   r = 1.0

Changing only the accumulator dtype doesn't help, since y_pred.square().sum() is still computed in float32 before being widened. The .to(dtype=...) in _update is what fixes it. A short comment there would stop someone from removing it later as redundant.

On master this gives 800000000.0 for five points at offset 1e4, with no NaN. That uses float32 inputs, which is what the docstring says to pass, so it might be worth a test too. The current test uses float64 inputs at offset 1e8.

The test does catch the bug. On this branch the file is 32 passed, 34 skipped. With pearson_correlation.py reverted to master and the test kept:

assert 1.0 ± 1.0e-06 == 0.0
  Obtained: 0.0
  Expected: 1.0 ± 1.0e-06

One question. The five accumulators are in _state_dict_all_req_keys, and _CollectionItem.load_value assigns the loaded tensor directly. So loading a checkpoint saved before this change brings back float32 tensors, and _update then casts the inputs to float32 through self._sum_of_y_preds.dtype. The fix would be undone without any warning. Casting in reset or on load would handle that. What do you think?

Smaller notes:

• The MPS fallback keeps float32 accumulators, so the bug remains there. That's fine as a limitation, but the new test is CPU-only while the rest of the file uses available_device.

ruff check and ruff format pass on both files.

LGTM.