Repository: Lightning-AI/torchmetrics

Issue:
Implement __getitem__ as "metric arithmetic"
## 🚀 Feature

Allow a user to define a new metric that takes an item out of an other metric.

Basically :

```python
iou = IoU(num_classes=2, reduction="none")
fg_iou = iou[0]
bg_iou = iou[1]
```

### Motivation

There are multiple metrics (like IoU and confusion matrix) that would benefit from the use of such a feature, and it is close to the mechanism of metric arithmetic.

### Pitch

This would only need to define 
```python
class Metric:
    ...
    def __getitem__(self, idx):
        return CompositionalMetric(lambda x: x[idx], self, None)
```

### Alternatives

The straightforward alternative is to use CompositionalMetric directly.


Discussion:
Hi! thanks for your contribution!, great first issue!
I think this great addition. @justusschock opinion?

Resolving patch:
diff --git a/torchmetrics/metric.py b/torchmetrics/metric.py
--- a/torchmetrics/metric.py
+++ b/torchmetrics/metric.py
@@ -126,10 +126,7 @@ def add_state(
             ValueError:
                 If ``dist_reduce_fx`` is not callable or one of ``"mean"``, ``"sum"``, ``"cat"``, ``None``.
         """
-        if (
-            not isinstance(default, Tensor) and not isinstance(default, list)  # noqa: W503
-            or (isinstance(default, list) and len(default) != 0)  # noqa: W503
-        ):
+        if (not isinstance(default, (Tensor, list)) or (isinstance(default, list) and default)):
             raise ValueError("state variable must be a tensor or any empty list (where you can append tensors)")
 
         if dist_reduce_fx == "sum":
@@ -304,7 +301,7 @@ def persistent(self, mode: bool = False):
         for key in self._persistent.keys():
             self._persistent[key] = mode
 
-    def state_dict(self, destination=None, prefix='', keep_vars=False):
+    def state_dict(self, destination=None, prefix="", keep_vars=False):
         destination = super().state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
         # Register metric states to be part of the state_dict
         for key in self._defaults.keys():
@@ -342,7 +339,7 @@ def __hash__(self):
             val = getattr(self, key)
             # Special case: allow list values, so long
             # as their elements are hashable
-            if hasattr(val, '__iter__') and not isinstance(val, Tensor):
+            if hasattr(val, "__iter__") and not isinstance(val, Tensor):
                 hash_vals.extend(val)
             else:
                 hash_vals.append(val)
@@ -449,6 +446,9 @@ def __neg__(self):
     def __pos__(self):
         return CompositionalMetric(torch.abs, self, None)
 
+    def __getitem__(self, idx):
+        return CompositionalMetric(lambda x: x[idx], self, None)
+
 
 def _neg(tensor: Tensor):
     return -torch.abs(tensor)
@@ -530,6 +530,6 @@ def persistent(self, mode: bool = False) -> None:
 
     def __repr__(self) -> str:
         _op_metrics = f"(\n  {self.op.__name__}(\n    {repr(self.metric_a)},\n    {repr(self.metric_b)}\n  )\n)"
-        repr_str = (self.__class__.__name__ + _op_metrics)
+        repr_str = self.__class__.__name__ + _op_metrics
 
         return repr_str
diff --git a/torchmetrics/wrappers/bootstrapping.py b/torchmetrics/wrappers/bootstrapping.py
--- a/torchmetrics/wrappers/bootstrapping.py
+++ b/torchmetrics/wrappers/bootstrapping.py
@@ -22,10 +22,7 @@
 from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_7
 
 
-def _bootstrap_sampler(
-    size: int,
-    sampling_strategy: str = 'poisson'
-) -> Tensor:
+def _bootstrap_sampler(size: int, sampling_strategy: str = 'poisson') -> Tensor:
     """ Resample a tensor along its first dimension with replacement
     Args:
         size: number of samples
@@ -38,14 +35,10 @@ def _bootstrap_sampler(
     """
     if sampling_strategy == 'poisson':
         p = torch.distributions.Poisson(1)
-        n = p.sample((size,))
+        n = p.sample((size, ))
         return torch.arange(size).repeat_interleave(n.long(), dim=0)
     elif sampling_strategy == 'multinomial':
-        idx = torch.multinomial(
-            torch.ones(size),
-            num_samples=size,
-            replacement=True
-        )
+        idx = torch.multinomial(torch.ones(size), num_samples=size, replacement=True)
         return idx
     raise ValueError('Unknown sampling strategy')
 


Test patch:
diff --git a/tests/bases/test_composition.py b/tests/bases/test_composition.py
--- a/tests/bases/test_composition.py
+++ b/tests/bases/test_composition.py
@@ -325,8 +325,10 @@ def test_metrics_rfloordiv(first_operand, expected_result):
     assert torch.allclose(expected_result, final_rfloordiv.compute())
 
 
-@pytest.mark.parametrize(["first_operand", "expected_result"],
-                         [pytest.param(tensor([2, 2, 2]), tensor(12), marks=pytest.mark.skipif(**_MARK_TORCH_MIN_1_4))])
+@pytest.mark.parametrize(
+    ["first_operand", "expected_result"],
+    [pytest.param(tensor([2, 2, 2]), tensor(12), marks=pytest.mark.skipif(**_MARK_TORCH_MIN_1_4))],
+)
 def test_metrics_rmatmul(first_operand, expected_result):
     second_operand = DummyMetric([2, 2, 2])
 
@@ -337,8 +339,10 @@ def test_metrics_rmatmul(first_operand, expected_result):
     assert torch.allclose(expected_result, final_rmatmul.compute())
 
 
-@pytest.mark.parametrize(["first_operand", "expected_result"],
-                         [pytest.param(tensor(2), tensor(2), marks=pytest.mark.skipif(**_MARK_TORCH_MIN_1_4))])
+@pytest.mark.parametrize(
+    ["first_operand", "expected_result"],
+    [pytest.param(tensor(2), tensor(2), marks=pytest.mark.skipif(**_MARK_TORCH_MIN_1_4))],
+)
 def test_metrics_rmod(first_operand, expected_result):
     second_operand = DummyMetric(5)
 
@@ -497,6 +501,18 @@ def test_metrics_pos():
     assert torch.allclose(tensor(1), final_pos.compute())
 
 
+@pytest.mark.parametrize(
+    ["value", "idx", "expected_result"],
+    [([1, 2, 3], 1, tensor(2)), ([[0, 1], [2, 3]], (1, 0), tensor(2)), ([[0, 1], [2, 3]], 1, tensor([2, 3]))],
+)
+def test_metrics_getitem(value, idx, expected_result):
+    first_metric = DummyMetric(value)
+
+    final_getitem = first_metric[idx]
+    assert isinstance(final_getitem, CompositionalMetric)
+    assert torch.allclose(expected_result, final_getitem.compute())
+
+
 def test_compositional_metrics_update():
 
     compos = DummyMetric(5) + DummyMetric(4)
diff --git a/tests/wrappers/test_bootstrapping.py b/tests/wrappers/test_bootstrapping.py
--- a/tests/wrappers/test_bootstrapping.py
+++ b/tests/wrappers/test_bootstrapping.py
@@ -32,6 +32,7 @@ class TestBootStrapper(BootStrapper):
     """ For testing purpose, we subclass the bootstrapper class so we can get the exact permutation
         the class is creating
     """
+
     def update(self, *args) -> None:
         self.out = []
         for idx in range(self.num_bootstraps):