Repository: Lightning-AI/torchmetrics

Issue:
`MeanMetric` with `nan_strategy in {"warn", "ignore"}` might not recover from NaNs
## 🐛 Bug

`MeanMetric` with `nan_strategy in {"warn", "ignore"}` does not recover from NaNs when example weights are given. 

I believe this is occuring because in [`MeanMetric.update()`](https://github.com/Lightning-AI/torchmetrics/blob/v1.0.0/src/torchmetrics/aggregation.py#L519-L520), the `value` and `weight` arguments are independently processed by removing any NaN entries. This can cause the shapes of the arguments to drift in a way that makes them no longer broadcastable. A fix could be to first broadcast the two tensors and then remove all indices at which `value` or `weight` are NaN.

### To Reproduce

```python
import torch
import torchmetrics

M = torchmetrics.MeanMetric(nan_strategy="ignore")  # same with "warn" 

x = torch.rand(5)
x[1] = torch.nan
w = torch.rand(5)

M.update(x, weight=w)
```
This raises an error:
```
Traceback (most recent call last):
  File "...", line 10, in <module>
    M.update(x, weight=w)
  File ".../python3.10/site-packages/torchmetrics/metric.py", line 400, in wrapped_func
    raise err
  File ".../python3.10/site-packages/torchmetrics/metric.py", line 390, in wrapped_func
    update(*args, **kwargs)
  File ".../python3.10/site-packages/torchmetrics/aggregation.py", line 402, in update
    weight = torch.broadcast_to(weight, value.shape)
RuntimeError: The expanded size of the tensor (4) must match the existing size (5) at non-singleton dimension 0.  Target sizes: [4].  Tensor sizes: [5]
```

### Expected behavior

An error should not be raised, as the program should continue and (optionally) provide a warning.

### Environment

- TorchMetrics version: `1.0.0`, installed with `pip`
- Python & PyTorch Version: `3.10.11` and `2.0.1`, respectively



Discussion:

Resolving patch:
diff --git a/src/torchmetrics/aggregation.py b/src/torchmetrics/aggregation.py
--- a/src/torchmetrics/aggregation.py
+++ b/src/torchmetrics/aggregation.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, List, Optional, Sequence, Union
+from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -69,25 +69,36 @@ def __init__(
         self.nan_strategy = nan_strategy
         self.add_state("value", default=default_value, dist_reduce_fx=fn)
 
-    def _cast_and_nan_check_input(self, x: Union[float, Tensor]) -> Tensor:
+    def _cast_and_nan_check_input(
+        self, x: Union[float, Tensor], weight: Optional[Union[float, Tensor]] = None
+    ) -> 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)
+        if weight is not None and not isinstance(weight, Tensor):
+            weight = torch.as_tensor(weight, dtype=torch.float32, device=self.device)
 
         nans = torch.isnan(x)
-        if nans.any():
+        if weight is not None:
+            nans_weight = torch.isnan(weight)
+        else:
+            nans_weight = torch.zeros_like(nans).bool()
+            weight = torch.ones_like(x)
+        if nans.any() or nans_weight.any():
             if self.nan_strategy == "error":
                 raise RuntimeError("Encounted `nan` values in tensor")
             if self.nan_strategy in ("ignore", "warn"):
                 if self.nan_strategy == "warn":
                     rank_zero_warn("Encounted `nan` values in tensor. Will be removed.", UserWarning)
-                x = x[~nans]
+                x = x[~(nans | nans_weight)]
+                weight = weight[~(nans | nans_weight)]
             else:
                 if not isinstance(self.nan_strategy, float):
                     raise ValueError(f"`nan_strategy` shall be float but you pass {self.nan_strategy}")
-                x[nans] = self.nan_strategy
+                x[nans | nans_weight] = self.nan_strategy
+                weight[nans | nans_weight] = self.nan_strategy
 
-        return x.float()
+        return x.float(), weight.float()
 
     def update(self, value: Union[float, Tensor]) -> None:
         """Overwrite in child class."""
@@ -153,7 +164,7 @@ def update(self, value: Union[float, Tensor]) -> None:
             value: Either a float or tensor containing data. Additional tensor
                 dimensions will be flattened
         """
-        value = self._cast_and_nan_check_input(value)
+        value, _ = self._cast_and_nan_check_input(value)
         if value.numel():  # make sure tensor not empty
             self.value = torch.max(self.value, torch.max(value))
 
@@ -253,7 +264,7 @@ def update(self, value: Union[float, Tensor]) -> None:
             value: Either a float or tensor containing data. Additional tensor
                 dimensions will be flattened
         """
-        value = self._cast_and_nan_check_input(value)
+        value, _ = self._cast_and_nan_check_input(value)
         if value.numel():  # make sure tensor not empty
             self.value = torch.min(self.value, torch.min(value))
 
@@ -351,7 +362,7 @@ def update(self, value: Union[float, Tensor]) -> None:
             value: Either a float or tensor containing data. Additional tensor
                 dimensions will be flattened
         """
-        value = self._cast_and_nan_check_input(value)
+        value, _ = self._cast_and_nan_check_input(value)
         if value.numel():
             self.value += value.sum()
 
@@ -445,7 +456,7 @@ def update(self, value: Union[float, Tensor]) -> None:
             value: Either a float or tensor containing data. Additional tensor
                 dimensions will be flattened
         """
-        value = self._cast_and_nan_check_input(value)
+        value, _ = self._cast_and_nan_check_input(value)
         if value.numel():
             self.value.append(value)
 
@@ -516,13 +527,16 @@ def update(self, value: Union[float, Tensor], weight: Union[float, Tensor] = 1.0
                 the shape of `value`. Default to `1.0` corresponding to simple
                 harmonic average.
         """
-        value = self._cast_and_nan_check_input(value)
-        weight = self._cast_and_nan_check_input(weight)
+        # broadcast weight to value shape
+        if not isinstance(value, Tensor):
+            value = torch.as_tensor(value, 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)
+        weight = torch.broadcast_to(weight, value.shape)
+        value, weight = self._cast_and_nan_check_input(value, weight)
 
         if value.numel() == 0:
             return
-        # broadcast weight to value shape
-        weight = torch.broadcast_to(weight, value.shape)
         self.value += (value * weight).sum()
         self.weight += weight.sum()
 


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
@@ -164,3 +164,26 @@ def test_mean_metric_broadcasting(weights, expected):
     avg = MeanMetric()
 
     assert avg(values, weights) == expected
+
+
+@pytest.mark.skipif(not hasattr(torch, "broadcast_to"), reason="PyTorch <1.8 does not have broadcast_to")
+@pytest.mark.parametrize("nan_strategy", ["ignore", "warn"])
+def test_mean_metric_broadcast(nan_strategy):
+    """Check that weights gets broadcasted correctly when Nans are present."""
+    metric = MeanMetric(nan_strategy=nan_strategy)
+
+    x = torch.arange(5).float()
+    x[1] = torch.tensor(float("nan"))
+    w = torch.arange(5).float()
+
+    metric.update(x, w)
+    res = metric.compute()
+    assert round(res.item(), 4) == 3.2222  # (0*0 + 2*2 + 3*3 + 4*4) / (0 + 2 + 3 + 4)
+
+    x = torch.arange(5).float()
+    w = torch.arange(5).float()
+    w[1] = torch.tensor(float("nan"))
+
+    metric.update(x, w)
+    res = metric.compute()
+    assert round(res.item(), 4) == 3.2222  # (0*0 + 2*2 + 3*3 + 4*4) / (0 + 2 + 3 + 4)