Repository: Lightning-AI/torchmetrics

Issue:
Explicitly initializing tensors in `float32` in MeanMetric goes against `torch.set_default_dtype`, leading to numerical errors
## 🐛 Bug

When working in `torch.float64` (set by `torch.set_default_dtype(torch.float64)`), the result of `MeanMetric` is slightly incorrect because, internaly, the `weight` and `value` tensors are initialized to `torch.float32` explicitly (see these [lines](https://github.com/Lightning-AI/torchmetrics/blob/c0d2d3ab7221f8804ba55baf3945637d68485cc1/src/torchmetrics/aggregation.py#L558)).

Naturally, if someone's working in float64 then the problem can't tolerate the slight mismatch.

### To Reproduce

```python
import torch
from torchmetrics.aggregation import MeanMetric

torch.set_default_dtype(torch.float64)

metric = MeanMetric()

values = torch.randn(10000)
metric.update(values)
result = metric.compute()

print(f"{result} = Result\n{values.mean()} = Actual mean")
```

The output of the snippet is:
```
0.018805044555664063 = Result
0.01880504383120193 = Actual mean
```

Notice that the two differ. Where as if we switch the default to `torch.float32`:
```
0.013675613328814507 = Result
0.013675613328814507 = Actual mean
```
they're exactly the same.

### Expected behavior

Idealy, the result of the MeanMetric (and possibly other metrics where this might have happened) should give the correct result irrespective of the default dtype.

### Environment

`pip freeze`:

```
torch==2.1.0+cu121
torchmetrics==1.2.1
```

Python 3.9

Windows 11



Discussion:
Hi! thanks for your contribution!, great first issue!
So honoring the default if it is set would help, right? mind sending a fix PR... :rabbit: 
Something like that.

I could try, would be my first time but I'll give it a go.
@viktor-ktorvi I had a bit of time to look at it and it should be fixed in PR #2366 
Hi, 

I just found the time to check this. I've upgraded and the issue still isn't fixed.

## What's wrong
The example I stated still doesn't function as expected.

## Why the test passed
The implemented [test](https://github.com/Lightning-AI/torchmetrics/blob/afae59e4d75ef2dfd0023e3a665e46a0a62833a5/tests/unittests/bases/test_aggregation.py#L213) passes because of the use of `torch.allclose`, however, the values need not be close, but exactly equal. There's no reason for the two to not be exactly equal if all the calculations are performed in the same dtype. 

Let me demonstrate why it's wrong:

```python
import torch
from torchmetrics.aggregation import MeanMetric

torch.set_default_dtype(torch.float64)

metric = MeanMetric()

values = torch.randn(10000)
metric.update(values)
result = metric.compute()

actual_mean = values.mean()

print(f"{result} = Result\n{actual_mean} = Actual mean")

print(f"\nAll close = {torch.allclose(result, actual_mean, atol=1e-12)}")
print(f"Exactly equal = {result == actual_mean}")
```

```
-0.0041637580871582034 = Result
-0.004163758971599815 = Actual mean

All close = True
Exactly equal = False
```

## Motivation
It might feel like I'm nitpicking, but these sorts of errors add up in complex problem formulations. For context, I'm working on approximating optimization problems in with ML, and in my particular case, when casting from `float64` to `float32` and recalculating the values, the equality and inequality constraints are no longer fulfilled and the objective function is off.

## How to fix

I've narrowed it down to the `_cast_and_nan_check_input` function, which gets called in `update` [line 564](https://github.com/Lightning-AI/torchmetrics/blob/afae59e4d75ef2dfd0023e3a665e46a0a62833a5/src/torchmetrics/aggregation.py#L564). At the end of `_cast_and_nan_check_input` ([line 104](https://github.com/Lightning-AI/torchmetrics/blob/afae59e4d75ef2dfd0023e3a665e46a0a62833a5/src/torchmetrics/aggregation.py#L104)), `x.float()` get's called, explicitly casting to `float32`, so that'd need changing.

Additionally, lots of these 
```python
if not isinstance(x, Tensor):
    x = torch.as_tensor(x, dtype=torch.float32, device=self.device)
if weight is not None and not isinstance(weight, Tensor):
    weight = torch.as_tensor(weight, dtype=torch.float32, device=self.device)
```
statements, where the dtype is explicitly called exist, e.g., [line 79](https://github.com/Lightning-AI/torchmetrics/blob/afae59e4d75ef2dfd0023e3a665e46a0a62833a5/src/torchmetrics/aggregation.py#L79) or [line 559](https://github.com/Lightning-AI/torchmetrics/blob/afae59e4d75ef2dfd0023e3a665e46a0a62833a5/src/torchmetrics/aggregation.py#L559). So, each of those would need to be replaced with `dtype=torch.get_default_dtype()`.

Finally, the [test](https://github.com/Lightning-AI/torchmetrics/blob/afae59e4d75ef2dfd0023e3a665e46a0a62833a5/tests/unittests/bases/test_aggregation.py#L213) needs to check for equality i.e., `result == compare_function(values)` instead of using `torch.allclose`.

Thanks for your time! @SkafteNicki @Borda 


Resolving patch:
diff --git a/src/torchmetrics/aggregation.py b/src/torchmetrics/aggregation.py
--- a/src/torchmetrics/aggregation.py
+++ b/src/torchmetrics/aggregation.py
@@ -77,9 +77,9 @@ def _cast_and_nan_check_input(
     ) -> Tuple[Tensor, Tensor]:
         """Convert input ``x`` to a tensor and check for Nans."""
         if not isinstance(x, Tensor):
-            x = torch.as_tensor(x, dtype=torch.float32, device=self.device)
+            x = torch.as_tensor(x, dtype=self.dtype, device=self.device)
         if weight is not None and not isinstance(weight, Tensor):
-            weight = torch.as_tensor(weight, dtype=torch.float32, device=self.device)
+            weight = torch.as_tensor(weight, dtype=self.dtype, device=self.device)
 
         nans = torch.isnan(x)
         if weight is not None:
@@ -101,7 +101,7 @@ def _cast_and_nan_check_input(
                 x[nans | nans_weight] = self.nan_strategy
                 weight[nans | nans_weight] = self.nan_strategy
 
-        return x.float(), weight.float()
+        return x.to(self.dtype), weight.to(self.dtype)
 
     def update(self, value: Union[float, Tensor]) -> None:
         """Overwrite in child class."""
@@ -557,9 +557,9 @@ def update(self, value: Union[float, Tensor], weight: Union[float, Tensor] = 1.0
         """
         # broadcast weight to value shape
         if not isinstance(value, Tensor):
-            value = torch.as_tensor(value, dtype=torch.float32, device=self.device)
+            value = torch.as_tensor(value, dtype=self.dtype, device=self.device)
         if weight is not None and not isinstance(weight, Tensor):
-            weight = torch.as_tensor(weight, dtype=torch.float32, device=self.device)
+            weight = torch.as_tensor(weight, dtype=self.dtype, device=self.device)
         weight = torch.broadcast_to(weight, value.shape)
         value, weight = self._cast_and_nan_check_input(value, weight)
 


Test patch:
diff --git a/tests/unittests/bases/test_aggregation.py b/tests/unittests/bases/test_aggregation.py
--- a/tests/unittests/bases/test_aggregation.py
+++ b/tests/unittests/bases/test_aggregation.py
@@ -214,10 +214,11 @@ def test_with_default_dtype(metric_class, compare_function):
     """Test that the metric works with a default dtype of float64."""
     torch.set_default_dtype(torch.float64)
     metric = metric_class()
+    assert metric.dtype == torch.float64
     values = torch.randn(10000)
     metric.update(values)
     result = metric.compute()
     assert result.dtype == torch.float64
     assert result.dtype == values.dtype
-    assert torch.allclose(result, compare_function(values), atol=1e-12)
+    assert result == compare_function(values)
     torch.set_default_dtype(torch.float32)
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
@@ -573,3 +573,19 @@ def test_update_properties(metric, method):
     m.reset()
     assert not m.update_called
     assert m.update_count == 0
+
+
+def test_dtype_property():
+    """Test that dtype property works as expected."""
+    metric = DummyMetricSum()
+    assert metric.dtype == torch.float32
+    metric.set_dtype(torch.float64)
+    assert metric.dtype == torch.float64
+
+    torch.set_default_dtype(torch.float64)
+    metric = DummyMetricSum()
+    assert metric.dtype == torch.float64
+    torch.set_default_dtype(torch.float32)
+    assert metric.dtype == torch.float64  # should not change after initialization
+    metric.set_dtype(torch.float32)
+    assert metric.dtype == torch.float32