Repository: Lightning-AI/torchmetrics

Issue:
dist_sync_fn seems dysfunctional
## 🐛 Bug

`dist_sync_fn` is a documented kwarg of class Metric and suggests to pass an alternative to `torch.distributed.all_gather`. I suppose, the intend is to allow use of distributed contexts other than `torch.distributed`, for which the default works fine. This use case fails due to the following issues:

- passed `dist_sync_fn` is never called because hardcoded test `jit_distributed_available()` returns False in any other sync context
- this test is assigned as default to kwarg `distributed_available` of `sync()` and `sync_context()`, but neither can be called by the user because all metrics already wrap their `compute()` method at init
- `sync_context` can be applied only once to a function, else the inner wrapper will raise TorchMetricsUserError "The Metric has already been synced."

### To Reproduce

Code for testing Accuracy and MetricCollection in a TPU context (colab): https://colab.research.google.com/drive/1MlxWSrkKKuZ3WSb9duf1c0MoO2A8jDAE?usp=sharing

#### Code sample

```py
acc = torchmetrics.Accuracy(dist_sync_fn=gather_fn).to(device)
acc.update(scores.detach(), labels)
acc_value = acc.compute()
```
### Expected behavior

Automatically sync and compute correct metrics.

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 0.10dev
- Python & PyTorch Version (e.g., 1.0):
- Any other relevant information such as OS (e.g., Linux): e.g. torch_xla.distributed

### Additional context

I found that with minimal changes in metrics.py, `dist_sync_fn` can actually be used to run torchmetrics on TPUs (torch_xla), like in the code sample above. I could send a PR (see fork in colab notebook).



Discussion:
Hi! thanks for your contribution!, great first issue!
@justusschock, mind having a look at it? :otter: 

Resolving patch:
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -69,6 +69,8 @@ class Metric(Module, ABC):
             - process_group: The process group on which the synchronization is called. Default is the world.
             - dist_sync_fn: function that performs the allgather option on the metric state. Default is an
                 custom implementation that calls ``torch.distributed.all_gather`` internally.
+            - distributed_available_fn: function that checks if the distributed backend is available.
+                Defaults to a check of ``torch.distributed.is_available()`` and ``torch.distributed.is_initialized()``.
             - sync_on_compute: If metric state should synchronize when ``compute`` is called. Default is ``True``-
     """
 
@@ -110,6 +112,8 @@ def __init__(
                 f"Expected keyword argument `dist_sync_fn` to be an callable function but got {self.dist_sync_fn}"
             )
 
+        self.distributed_available_fn = kwargs.pop("distributed_available_fn", jit_distributed_available)
+
         self.sync_on_compute = kwargs.pop("sync_on_compute", True)
         if not isinstance(self.sync_on_compute, bool):
             raise ValueError(
@@ -421,7 +425,7 @@ def sync(
         dist_sync_fn: Optional[Callable] = None,
         process_group: Optional[Any] = None,
         should_sync: bool = True,
-        distributed_available: Optional[Callable] = jit_distributed_available,
+        distributed_available: Optional[Callable] = None,
     ) -> None:
         """Sync function for manually controlling when metrics states should be synced across processes.
 
@@ -437,6 +441,9 @@ def sync(
         if self._is_synced and should_sync:
             raise TorchMetricsUserError("The Metric has already been synced.")
 
+        if distributed_available is None and self.distributed_available_fn is not None:
+            distributed_available = self.distributed_available_fn
+
         is_distributed = distributed_available() if callable(distributed_available) else None
 
         if not should_sync or not is_distributed:
@@ -481,7 +488,7 @@ def sync_context(
         process_group: Optional[Any] = None,
         should_sync: bool = True,
         should_unsync: bool = True,
-        distributed_available: Optional[Callable] = jit_distributed_available,
+        distributed_available: Optional[Callable] = None,
     ) -> Generator:
         """Context manager to synchronize the states between processes when running in a distributed setting and
         restore the local cache states after yielding.


Test patch:
diff --git a/tests/unittests/bases/test_metric.py b/tests/unittests/bases/test_metric.py
--- a/tests/unittests/bases/test_metric.py
+++ b/tests/unittests/bases/test_metric.py
@@ -14,6 +14,7 @@
 import os
 import pickle
 from collections import OrderedDict
+from unittest.mock import Mock
 
 import cloudpickle
 import numpy as np
@@ -23,7 +24,7 @@
 from torch import Tensor, tensor
 from torch.nn import Module
 
-from torchmetrics import PearsonCorrCoef
+from torchmetrics import Accuracy, PearsonCorrCoef
 from unittests.helpers import seed_all
 from unittests.helpers.testers import DummyListMetric, DummyMetric, DummyMetricMultiOutput, DummyMetricSum
 from unittests.helpers.utilities import no_warning_call
@@ -449,3 +450,17 @@ def forward(self, *args, **kwargs):
         match="Torchmetrics v0.9 introduced a new argument class property called.*",
     ):
         UnsetProperty()
+
+
+def test_custom_availability_check_and_sync_fn():
+    dummy_availability_check = Mock(return_value=True)
+    dummy_dist_sync_fn = Mock(wraps=lambda x, group: [x])
+    acc = Accuracy(dist_sync_fn=dummy_dist_sync_fn, distributed_available_fn=dummy_availability_check)
+
+    acc.update(torch.tensor([[1], [1], [1], [1]]), torch.tensor([[1], [1], [1], [1]]))
+    dummy_dist_sync_fn.assert_not_called()
+    dummy_availability_check.assert_not_called()
+
+    acc.compute()
+    dummy_availability_check.assert_called_once()
+    assert dummy_dist_sync_fn.call_count == 4  # tp, fp, tn, fn