Repository: Lightning-AI/torchmetrics

Issue:
PearsonCorrCoef UnboundLocalError using ddp
I want to use PearsonCorrCoef in my pytorch lightning code with ddp strategy.
My num_outputs > 1
In x_step_end, I use CorrCoef.update(preds, targets)
In x_epoch_end, I use CorrCoef.compute().mean()
Always received

```
  File "/home/xx/Documents/anconda3/envs/xx/lib/python3.8/site-packages/torchmetrics/metric.py", line 531, in wrapped_func
    value = compute(*args, **kwargs)
  File "/home/xx/Documents/anconda3/envs/xx/lib/python3.8/site-packages/torchmetrics/regression/pearson.py", line 152, in compute
    _, _, var_x, var_y, corr_xy, n_total = _final_aggregation(
  File "/home/xx/Documents/anconda3/envs/xx/lib/python3.8/site-packages/torchmetrics/regression/pearson.py", line 63, in _final_aggregation
    return mean_x, mean_y, var_x, var_y, corr_xy, nb
UnboundLocalError: local variable 'mean_x' referenced before assignment
```

But when I change to dp strategy, it's fine.



Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/src/torchmetrics/regression/pearson.py b/src/torchmetrics/regression/pearson.py
--- a/src/torchmetrics/regression/pearson.py
+++ b/src/torchmetrics/regression/pearson.py
@@ -32,7 +32,8 @@ def _final_aggregation(
 
     Formula taken from here: `Aggregate the statistics from multiple devices`_
     """
-    # assert len(means_x) > 1 and len(means_y) > 1 and len(vars_x) > 1 and len(vars_y) > 1 and len(corrs_xy) > 1
+    if len(means_x) == 1:
+        return means_x[0], means_y[0], vars_x[0], vars_y[0], corrs_xy[0], nbs[0]
     mx1, my1, vx1, vy1, cxy1, n1 = means_x[0], means_y[0], vars_x[0], vars_y[0], corrs_xy[0], nbs[0]
     for i in range(1, len(means_x)):
         mx2, my2, vx2, vy2, cxy2, n2 = means_x[i], means_y[i], vars_x[i], vars_y[i], corrs_xy[i], nbs[i]


Test patch:
diff --git a/tests/unittests/regression/test_pearson.py b/tests/unittests/regression/test_pearson.py
--- a/tests/unittests/regression/test_pearson.py
+++ b/tests/unittests/regression/test_pearson.py
@@ -19,7 +19,7 @@
 from scipy.stats import pearsonr
 
 from torchmetrics.functional.regression.pearson import pearson_corrcoef
-from torchmetrics.regression.pearson import PearsonCorrCoef
+from torchmetrics.regression.pearson import PearsonCorrCoef, _final_aggregation
 from unittests import BATCH_SIZE, EXTRA_DIM, NUM_BATCHES
 from unittests.helpers import seed_all
 from unittests.helpers.testers import MetricTester
@@ -121,3 +121,12 @@ def test_error_on_different_shape():
     metric = PearsonCorrCoef(num_outputs=2)
     with pytest.raises(ValueError, match="Expected argument `num_outputs` to match the second dimension of input.*"):
         metric(torch.randn(100, 5), torch.randn(100, 5))
+
+
+@pytest.mark.parametrize("shapes", [(5,), (1, 5), (2, 5)])
+def test_final_aggregation_function(shapes):
+    """Test that final aggregation function can take various shapes of input."""
+    input_fn = lambda: torch.rand(shapes)
+    output = _final_aggregation(input_fn(), input_fn(), input_fn(), input_fn(), input_fn(), torch.randint(10, shapes))
+    assert all(isinstance(out, torch.Tensor) for out in output)
+    assert all(out.ndim == input_fn().ndim - 1 for out in output)