Repository: Lightning-AI/torchmetrics

Issue:
BootStrapper.update/forward don't process kwargs
## 🐛 Bug

Both `BootStrapper.update` and `BootStrapper.forward` fail when called with keyword arguments. This is because a dictionary is converted to a list with `list(d)` which produces a list of keys instead of a list of values

### To Reproduce

Defining a BootStrapper on a metric works when called arguments but fails when called with a keyword argument.

<details>
<summary>
Calling a bootstrapped metric with regular arguments works as expected, but using keyword argument raises an error because the inferred `size` is a string.

```python
import torch
import torchmetrics as tm

x=torch.uniform(100)
y=x+torch.randn_like(x)

ae = tm.MeanAbsoluteError()
print(ae(x,y)) # > tensor(...)
print(ae(preds=x,target=y) # > tensor(...)

bootstrapped_ae = tm.BootStrapper(ae)

print(bootstrapped_ae(x,y)) # > {'mean': tensor(...), 'std': tensor(...)}

print(bootstrapped_ae(preds=x,target=y)) # raises an error
```

**Culprit:** the following statement produces a list of keys (strings) instead of a list of `Sizes`

```python
# torchmetrics/wrappers/bootstrapping.py:131
        kwargs_sizes = list(apply_to_collection(kwargs, Tensor, len))
```

</summary>

Stacktrace:
```
File ~/.pyenv/versions/3.10.13/envs/br-div-dhnew/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File ~/.pyenv/versions/3.10.13/envs/br-div-dhnew/lib/python3.10/site-packages/torchmetrics/wrappers/bootstrapping.py:168, in BootStrapper.forward(self, *args, **kwargs)
    166 def forward(self, *args: Any, **kwargs: Any) -> Any:
    167     """Use the original forward method of the base metric class."""
--> 168     return super(WrapperMetric, self).forward(*args, **kwargs)

File ~/.pyenv/versions/3.10.13/envs/br-div-dhnew/lib/python3.10/site-packages/torchmetrics/metric.py:301, in Metric.forward(self, *args, **kwargs)
    295     raise TorchMetricsUserError(
    296         "The Metric shouldn't be synced when performing ``forward``. "
    297         "HINT: Did you forget to call ``unsync`` ?."
    298     )
    300 if self.full_state_update or self.full_state_update is None or self.dist_sync_on_step:
--> 301     self._forward_cache = self._forward_full_state_update(*args, **kwargs)
    302 else:
    303     self._forward_cache = self._forward_reduce_state_update(*args, **kwargs)

File ~/.pyenv/versions/3.10.13/envs/br-div-dhnew/lib/python3.10/site-packages/torchmetrics/metric.py:316, in Metric._forward_full_state_update(self, *args, **kwargs)
    308 """Forward computation using two calls to `update`.
    309
    310 Doing this secures that metrics that need access to the full metric state during `update` works as expected.
   (...)
    313
    314 """
    315 # global accumulation
--> 316 self.update(*args, **kwargs)
    317 _update_count = self._update_count
    319 self._to_sync = self.dist_sync_on_step

File ~/.pyenv/versions/3.10.13/envs/br-div-dhnew/lib/python3.10/site-packages/torchmetrics/wrappers/bootstrapping.py:140, in BootStrapper.update(self, *args, **kwargs)
    137     raise ValueError("None of the input contained tensors, so could not determine the sampling size")
    139 for idx in range(self.num_bootstraps):
--> 140     sample_idx = _bootstrap_sampler(size, sampling_strategy=self.sampling_strategy).to(self.device)
    141     if sample_idx.numel() == 0:
    142         continue

File ~/.pyenv/versions/3.10.13/envs/br-div-dhnew/lib/python3.10/site-packages/torchmetrics/wrappers/bootstrapping.py:47, in _bootstrap_sampler(size, sampling_strategy)
     45 if sampling_strategy == "poisson":
     46     p = torch.distributions.Poisson(1)
---> 47     n = p.sample((size,))
     48     return torch.arange(size).repeat_interleave(n.long(), dim=0)
     49 if sampling_strategy == "multinomial":

File ~/.pyenv/versions/3.10.13/envs/br-div-dhnew/lib/python3.10/site-packages/torch/distributions/poisson.py:61, in Poisson.sample(self, sample_shape)
     60 def sample(self, sample_shape=torch.Size()):
---> 61     shape = self._extended_shape(sample_shape)
     62     with torch.no_grad():
     63         return torch.poisson(self.rate.expand(shape))

File ~/.pyenv/versions/3.10.13/envs/br-div-dhnew/lib/python3.10/site-packages/torch/distributions/distribution.py:259, in Distribution._extended_shape(self, sample_shape)
    249 """
    250 Returns the size of the sample returned by the distribution, given
    251 a `sample_shape`. Note, that the batch and event shapes of a distribution
   (...)
    256     sample_shape (torch.Size): the size of the sample to be drawn.
    257 """
    258 if not isinstance(sample_shape, torch.Size):
--> 259     sample_shape = torch.Size(sample_shape)
    260 return torch.Size(sample_shape + self._batch_shape + self._event_shape)

TypeError: torch.Size() takes an iterable of 'int' (item 0 is 'str')
```

</details>

### Expected behavior
The forward and update methods of `BootStrapper` should behave like those of the wrapped metric

### Environment

- TorchMetrics version: 1.3.1
- Python & PyTorch Version: 3.10.13 & 2.0.1
- Any other relevant information such as OS: OSX



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

Resolving patch:
diff --git a/src/torchmetrics/wrappers/bootstrapping.py b/src/torchmetrics/wrappers/bootstrapping.py
--- a/src/torchmetrics/wrappers/bootstrapping.py
+++ b/src/torchmetrics/wrappers/bootstrapping.py
@@ -129,11 +129,11 @@ def update(self, *args: Any, **kwargs: Any) -> None:
 
         """
         args_sizes = apply_to_collection(args, Tensor, len)
-        kwargs_sizes = list(apply_to_collection(kwargs, Tensor, len))
+        kwargs_sizes = apply_to_collection(kwargs, Tensor, len)
         if len(args_sizes) > 0:
             size = args_sizes[0]
         elif len(kwargs_sizes) > 0:
-            size = kwargs_sizes[0]
+            size = next(iter(kwargs_sizes.values()))
         else:
             raise ValueError("None of the input contained tensors, so could not determine the sampling size")
 


Test patch:
diff --git a/tests/unittests/wrappers/test_bootstrapping.py b/tests/unittests/wrappers/test_bootstrapping.py
--- a/tests/unittests/wrappers/test_bootstrapping.py
+++ b/tests/unittests/wrappers/test_bootstrapping.py
@@ -22,7 +22,7 @@
 from sklearn.metrics import mean_squared_error, precision_score, recall_score
 from torch import Tensor
 from torchmetrics.classification import MulticlassF1Score, MulticlassPrecision, MulticlassRecall
-from torchmetrics.regression import MeanSquaredError
+from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError
 from torchmetrics.wrappers.bootstrapping import BootStrapper, _bootstrap_sampler
 
 from unittests._helpers import seed_all
@@ -140,3 +140,23 @@ def test_low_sample_amount(sampling_strategy):
         MulticlassF1Score(num_classes=3, average=None), num_bootstraps=20, sampling_strategy=sampling_strategy
     )
     assert bootstrap_f1(preds, target)  # does not work
+
+
+def test_args_and_kwargs_works():
+    """Test that metric works with both args and kwargs and mix.
+
+    See issue: https://github.com/Lightning-AI/torchmetrics/issues/2450
+
+    """
+    x = torch.rand(100)
+    y = x + torch.randn_like(x)
+    ae = MeanAbsoluteError()
+    assert ae(x, y) == ae(preds=x, target=y)
+
+    bootstrapped_ae = BootStrapper(ae)
+    res1 = bootstrapped_ae(x, y)
+    res2 = bootstrapped_ae(x, target=y)
+    res3 = bootstrapped_ae(preds=x, target=y)
+
+    assert (res1["mean"].shape == res2["mean"].shape) & (res2["mean"].shape == res3["mean"].shape)
+    assert (res1["std"].shape == res2["std"].shape) & (res2["mean"].shape == res3["std"].shape)