Repository: Lightning-AI/torchmetrics

Issue:
State Changes from List to Tensor After Calling Compute()
## 🐛 Bug

I was logging a custom Metric result using the logger provided by LightningModule. The codes runs on 2 GPUs with ddp. While a list state whose `dist_reduce_fx` is "cat" was synchronized by calling `compute()` , I realized the behavior in `training_step()` and `validation_step()` were different.

The state in `training_step()` where on_step=True and on_epoch=False (default setting for training) after calling `compute()` was always a list. However, in `validation_step()` where on_step=False and on_epoch=True (default setting for validation), the state became a Tensor. Such behavior is not explained in the doc.

### To Reproduce

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

<details>
  <summary>Metric and Logging Results</summary>

**Codes**
```python
import logging

import torch
from torchmetrics import Metric
from torchmetrics.functional.retrieval import retrieval_normalized_dcg


class MyMetric(Metric):
    def __init__(self, top_k):
        super().__init__()
        self.top_k = top_k
        self.add_state("state_1", default=[], dist_reduce_fx="cat")

    def update(self, preds: Tensor, target: Tensor):
        assert preds.shape == target.shape
        self.state_1 += [self._metric(p, t) for p, t in zip(preds, target)]

    def compute(self):
        if isinstance(self.state_1, list):
            logging.warning(f"self.state_1.size(): {len(self.state_1)}")
            return torch.stack(self.state_1).mean()
        logging.warning(f"self.state_1.size(): {self.state_1.size()}")
        return self.state_1.mean()

    def _metric(self, preds: Tensor, target: Tensor):
        return retrieval_normalized_dcg(preds, target, k=self.top_k).float()
```

**Log Outputs**
Training:
`WARNING:self.state_1.size(): 40`

Validation:
`WARNING:self.state_1.size(): torch.Size([80])`

</details>

### Expected behavior

The state is always what it is defined in add_state(), i.e., a list.

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source):
0.10.3 (conda-forge)
- Python & PyTorch Version (e.g., 1.0):
3.10.12
- Any other relevant information such as OS (e.g., Linux):
pytorch: 1.13.1
pytorch-lightning: 1.9.4
Linux: 5.15.0-78-generic

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
@donglihe-hub, thanks for raising this issue.
I definitely see the problem that you have pointed out, the core problem with changing the behaviour being backwards compatibility. This behaviour has been present since the start of torchmetrics and we have internally really not considered it a problem. The solution in our own metrics are just to call the `dim_zero_cat` function in compute:
```python
from torchmetrics.metric import Metric
from torchmetrics.utilities.data import dim_zero_cat
class MyMetric(Metric):
     ...
     def compute(self):
         state_1 = dim_zero_cat(self.state_1)  # the returned value will always be a tensor
```
which secures that everything is a tensor before doing any computation.

cc: @justusschock opinions on this issue? is it something we should change?




> @donglihe-hub, thanks for raising this issue. I definitely see the problem that you have pointed out, the core problem with changing the behaviour being backwards compatibility. This behaviour has been present since the start of torchmetrics and we have internally really not considered it a problem. The solution in our own metrics are just to call the `dim_zero_cat` function in compute:
> 
> ```python
> from torchmetrics.metric import Metric
> from torchmetrics.utilities.data import dim_zero_cat
> class MyMetric(Metric):
>      ...
>      def compute(self):
>          state_1 = dim_zero_cat(self.state_1)  # the returned value will always be a tensor
> ```
> 
> which secures that everything is a tensor before doing any computation.
> 
> cc: @justusschock opinions on this issue? is it something we should change?

Thank you for your answer. I think maybe it should be mentioned in [IMPLEMENTING A METRIC](https://torchmetrics.readthedocs.io/en/stable/pages/implement.html)? Otherwise users have no way to know that before they actually running their codes and met the "bug".
@donglihe-hub you are completely right. 
That page have not been updated in a long time a definite need an overhaul. Let me try to figure out some more examples and hopefully we can also figure out the issue you have reported in #2022.

Resolving patch:
diff --git a/src/torchmetrics/utilities/__init__.py b/src/torchmetrics/utilities/__init__.py
--- a/src/torchmetrics/utilities/__init__.py
+++ b/src/torchmetrics/utilities/__init__.py
@@ -12,6 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from torchmetrics.utilities.checks import check_forward_full_state_property
+from torchmetrics.utilities.data import (
+    dim_zero_cat,
+    dim_zero_max,
+    dim_zero_mean,
+    dim_zero_min,
+    dim_zero_sum,
+)
 from torchmetrics.utilities.distributed import class_reduce, reduce
 from torchmetrics.utilities.prints import rank_zero_debug, rank_zero_info, rank_zero_warn
 
@@ -22,4 +29,9 @@
     "rank_zero_debug",
     "rank_zero_info",
     "rank_zero_warn",
+    "dim_zero_cat",
+    "dim_zero_max",
+    "dim_zero_mean",
+    "dim_zero_min",
+    "dim_zero_sum",
 ]
diff --git a/src/torchmetrics/utilities/distributed.py b/src/torchmetrics/utilities/distributed.py
--- a/src/torchmetrics/utilities/distributed.py
+++ b/src/torchmetrics/utilities/distributed.py
@@ -105,8 +105,7 @@ def gather_all_tensors(result: Tensor, group: Optional[Any] = None) -> List[Tens
         group: the process group to gather results from. Defaults to all processes (world)
 
     Return:
-        gathered_result: list with size equal to the process group where
-            ``gathered_result[i]`` corresponds to result tensor from process ``i``
+        list with size equal to the process group where element i corresponds to result tensor from process i
 
     """
     if group is None:


Test patch: