We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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. -->


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/torchmetrics/utilities/__init__.py`
Content:
```
1 # Copyright The Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from torchmetrics.utilities.checks import check_forward_full_state_property
15 from torchmetrics.utilities.distributed import class_reduce, reduce
16 from torchmetrics.utilities.prints import rank_zero_debug, rank_zero_info, rank_zero_warn
17 
18 __all__ = [
19     "check_forward_full_state_property",
20     "class_reduce",
21     "reduce",
22     "rank_zero_debug",
23     "rank_zero_info",
24     "rank_zero_warn",
25 ]
26
```

Path: `src/torchmetrics/utilities/distributed.py`
Content:
```
1 # Copyright The Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from typing import Any, List, Optional
15 
16 import torch
17 from torch import Tensor
18 from torch.nn import functional as F  # noqa: N812
19 from typing_extensions import Literal
20 
21 
22 def reduce(x: Tensor, reduction: Literal["elementwise_mean", "sum", "none", None]) -> Tensor:
23     """Reduces a given tensor by a given reduction method.
24 
25     Args:
26         x: the tensor, which shall be reduced
27         reduction:  a string specifying the reduction method ('elementwise_mean', 'none', 'sum')
28 
29     Return:
30         reduced Tensor
31 
32     Raise:
33         ValueError if an invalid reduction parameter was given
34 
35     """
36     if reduction == "elementwise_mean":
37         return torch.mean(x)
38     if reduction == "none" or reduction is None:
39         return x
40     if reduction == "sum":
41         return torch.sum(x)
42     raise ValueError("Reduction parameter unknown.")
43 
44 
45 def class_reduce(
46     num: Tensor,
47     denom: Tensor,
48     weights: Tensor,
49     class_reduction: Literal["micro", "macro", "weighted", "none", None] = "none",
50 ) -> Tensor:
51     """Reduce classification metrics of the form ``num / denom * weights``.
52 
53     For example for calculating standard accuracy the num would be number of true positives per class, denom would be
54     the support per class, and weights would be a tensor of 1s.
55 
56     Args:
57         num: numerator tensor
58         denom: denominator tensor
59         weights: weights for each class
60         class_reduction: reduction method for multiclass problems:
61 
62             - ``'micro'``: calculate metrics globally (default)
63             - ``'macro'``: calculate metrics for each label, and find their unweighted mean.
64             - ``'weighted'``: calculate metrics for each label, and find their weighted mean.
65             - ``'none'`` or ``None``: returns calculated metric per class
66 
67     Raises:
68         ValueError:
69             If ``class_reduction`` is none of ``"micro"``, ``"macro"``, ``"weighted"``, ``"none"`` or ``None``.
70 
71     """
72     valid_reduction = ("micro", "macro", "weighted", "none", None)
73     fraction = torch.sum(num) / torch.sum(denom) if class_reduction == "micro" else num / denom
74 
75     # We need to take care of instances where the denom can be 0
76     # for some (or all) classes which will produce nans
77     fraction[fraction != fraction] = 0
78 
79     if class_reduction == "micro":
80         return fraction
81     if class_reduction == "macro":
82         return torch.mean(fraction)
83     if class_reduction == "weighted":
84         return torch.sum(fraction * (weights.float() / torch.sum(weights)))
85     if class_reduction == "none" or class_reduction is None:
86         return fraction
87 
88     raise ValueError(f"Reduction parameter {class_reduction} unknown. Choose between one of these: {valid_reduction}")
89 
90 
91 def _simple_gather_all_tensors(result: Tensor, group: Any, world_size: int) -> List[Tensor]:
92     gathered_result = [torch.zeros_like(result) for _ in range(world_size)]
93     torch.distributed.all_gather(gathered_result, result, group)
94     return gathered_result
95 
96 
97 def gather_all_tensors(result: Tensor, group: Optional[Any] = None) -> List[Tensor]:
98     """Gather all tensors from several ddp processes onto a list that is broadcasted to all processes.
99 
100     Works on tensors that have the same number of dimensions, but where each dimension may differ. In this case
101     tensors are padded, gathered and then trimmed to secure equal workload for all processes.
102 
103     Args:
104         result: the value to sync
105         group: the process group to gather results from. Defaults to all processes (world)
106 
107     Return:
108         gathered_result: list with size equal to the process group where
109             ``gathered_result[i]`` corresponds to result tensor from process ``i``
110 
111     """
112     if group is None:
113         group = torch.distributed.group.WORLD
114 
115     # convert tensors to contiguous format
116     result = result.contiguous()
117 
118     world_size = torch.distributed.get_world_size(group)
119     torch.distributed.barrier(group=group)
120 
121     # if the tensor is scalar, things are easy
122     if result.ndim == 0:
123         return _simple_gather_all_tensors(result, group, world_size)
124 
125     # 1. Gather sizes of all tensors
126     local_size = torch.tensor(result.shape, device=result.device)
127     local_sizes = [torch.zeros_like(local_size) for _ in range(world_size)]
128     torch.distributed.all_gather(local_sizes, local_size, group=group)
129     max_size = torch.stack(local_sizes).max(dim=0).values
130     all_sizes_equal = all(all(ls == max_size) for ls in local_sizes)
131 
132     # 2. If shapes are all the same, then do a simple gather:
133     if all_sizes_equal:
134         return _simple_gather_all_tensors(result, group, world_size)
135 
136     # 3. If not, we need to pad each local tensor to maximum size, gather and then truncate
137     pad_dims = []
138     pad_by = (max_size - local_size).detach().cpu()
139     for val in reversed(pad_by):
140         pad_dims.append(0)
141         pad_dims.append(val.item())
142     result_padded = F.pad(result, pad_dims)
143     gathered_result = [torch.zeros_like(result_padded) for _ in range(world_size)]
144     torch.distributed.all_gather(gathered_result, result_padded, group)
145     for idx, item_size in enumerate(local_sizes):
146         slice_param = [slice(dim_size) for dim_size in item_size]
147         gathered_result[idx] = gathered_result[idx][slice_param]
148     return gathered_result
149
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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:

{"golden_diff": "diff --git a/src/torchmetrics/utilities/__init__.py b/src/torchmetrics/utilities/__init__.py\n--- a/src/torchmetrics/utilities/__init__.py\n+++ b/src/torchmetrics/utilities/__init__.py\n@@ -12,6 +12,13 @@\n # See the License for the specific language governing permissions and\n # limitations under the License.\n from torchmetrics.utilities.checks import check_forward_full_state_property\n+from torchmetrics.utilities.data import (\n+    dim_zero_cat,\n+    dim_zero_max,\n+    dim_zero_mean,\n+    dim_zero_min,\n+    dim_zero_sum,\n+)\n from torchmetrics.utilities.distributed import class_reduce, reduce\n from torchmetrics.utilities.prints import rank_zero_debug, rank_zero_info, rank_zero_warn\n \n@@ -22,4 +29,9 @@\n     \"rank_zero_debug\",\n     \"rank_zero_info\",\n     \"rank_zero_warn\",\n+    \"dim_zero_cat\",\n+    \"dim_zero_max\",\n+    \"dim_zero_mean\",\n+    \"dim_zero_min\",\n+    \"dim_zero_sum\",\n ]\ndiff --git a/src/torchmetrics/utilities/distributed.py b/src/torchmetrics/utilities/distributed.py\n--- a/src/torchmetrics/utilities/distributed.py\n+++ b/src/torchmetrics/utilities/distributed.py\n@@ -105,8 +105,7 @@\n         group: the process group to gather results from. Defaults to all processes (world)\n \n     Return:\n-        gathered_result: list with size equal to the process group where\n-            ``gathered_result[i]`` corresponds to result tensor from process ``i``\n+        list with size equal to the process group where element i corresponds to result tensor from process i\n \n     \"\"\"\n     if group is None:\n", "issue": "State Changes from List to Tensor After Calling Compute()\n## \ud83d\udc1b Bug\r\n\r\nI 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.\r\n\r\nThe 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.\r\n\r\n### To Reproduce\r\n\r\n<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->\r\n\r\n<details>\r\n  <summary>Metric and Logging Results</summary>\r\n\r\n**Codes**\r\n```python\r\nimport logging\r\n\r\nimport torch\r\nfrom torchmetrics import Metric\r\nfrom torchmetrics.functional.retrieval import retrieval_normalized_dcg\r\n\r\n\r\nclass MyMetric(Metric):\r\n    def __init__(self, top_k):\r\n        super().__init__()\r\n        self.top_k = top_k\r\n        self.add_state(\"state_1\", default=[], dist_reduce_fx=\"cat\")\r\n\r\n    def update(self, preds: Tensor, target: Tensor):\r\n        assert preds.shape == target.shape\r\n        self.state_1 += [self._metric(p, t) for p, t in zip(preds, target)]\r\n\r\n    def compute(self):\r\n        if isinstance(self.state_1, list):\r\n            logging.warning(f\"self.state_1.size(): {len(self.state_1)}\")\r\n            return torch.stack(self.state_1).mean()\r\n        logging.warning(f\"self.state_1.size(): {self.state_1.size()}\")\r\n        return self.state_1.mean()\r\n\r\n    def _metric(self, preds: Tensor, target: Tensor):\r\n        return retrieval_normalized_dcg(preds, target, k=self.top_k).float()\r\n```\r\n\r\n**Log Outputs**\r\nTraining:\r\n`WARNING:self.state_1.size(): 40`\r\n\r\nValidation:\r\n`WARNING:self.state_1.size(): torch.Size([80])`\r\n\r\n</details>\r\n\r\n### Expected behavior\r\n\r\nThe state is always what it is defined in add_state(), i.e., a list.\r\n\r\n### Environment\r\n\r\n- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source):\r\n0.10.3 (conda-forge)\r\n- Python & PyTorch Version (e.g., 1.0):\r\n3.10.12\r\n- Any other relevant information such as OS (e.g., Linux):\r\npytorch: 1.13.1\r\npytorch-lightning: 1.9.4\r\nLinux: 5.15.0-78-generic\r\n\r\n### Additional context\r\n\r\n<!-- Add any other context about the problem here. -->\r\n\n", "before_files": [{"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom torchmetrics.utilities.checks import check_forward_full_state_property\nfrom torchmetrics.utilities.distributed import class_reduce, reduce\nfrom torchmetrics.utilities.prints import rank_zero_debug, rank_zero_info, rank_zero_warn\n\n__all__ = [\n    \"check_forward_full_state_property\",\n    \"class_reduce\",\n    \"reduce\",\n    \"rank_zero_debug\",\n    \"rank_zero_info\",\n    \"rank_zero_warn\",\n]\n", "path": "src/torchmetrics/utilities/__init__.py"}, {"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, List, Optional\n\nimport torch\nfrom torch import Tensor\nfrom torch.nn import functional as F  # noqa: N812\nfrom typing_extensions import Literal\n\n\ndef reduce(x: Tensor, reduction: Literal[\"elementwise_mean\", \"sum\", \"none\", None]) -> Tensor:\n    \"\"\"Reduces a given tensor by a given reduction method.\n\n    Args:\n        x: the tensor, which shall be reduced\n        reduction:  a string specifying the reduction method ('elementwise_mean', 'none', 'sum')\n\n    Return:\n        reduced Tensor\n\n    Raise:\n        ValueError if an invalid reduction parameter was given\n\n    \"\"\"\n    if reduction == \"elementwise_mean\":\n        return torch.mean(x)\n    if reduction == \"none\" or reduction is None:\n        return x\n    if reduction == \"sum\":\n        return torch.sum(x)\n    raise ValueError(\"Reduction parameter unknown.\")\n\n\ndef class_reduce(\n    num: Tensor,\n    denom: Tensor,\n    weights: Tensor,\n    class_reduction: Literal[\"micro\", \"macro\", \"weighted\", \"none\", None] = \"none\",\n) -> Tensor:\n    \"\"\"Reduce classification metrics of the form ``num / denom * weights``.\n\n    For example for calculating standard accuracy the num would be number of true positives per class, denom would be\n    the support per class, and weights would be a tensor of 1s.\n\n    Args:\n        num: numerator tensor\n        denom: denominator tensor\n        weights: weights for each class\n        class_reduction: reduction method for multiclass problems:\n\n            - ``'micro'``: calculate metrics globally (default)\n            - ``'macro'``: calculate metrics for each label, and find their unweighted mean.\n            - ``'weighted'``: calculate metrics for each label, and find their weighted mean.\n            - ``'none'`` or ``None``: returns calculated metric per class\n\n    Raises:\n        ValueError:\n            If ``class_reduction`` is none of ``\"micro\"``, ``\"macro\"``, ``\"weighted\"``, ``\"none\"`` or ``None``.\n\n    \"\"\"\n    valid_reduction = (\"micro\", \"macro\", \"weighted\", \"none\", None)\n    fraction = torch.sum(num) / torch.sum(denom) if class_reduction == \"micro\" else num / denom\n\n    # We need to take care of instances where the denom can be 0\n    # for some (or all) classes which will produce nans\n    fraction[fraction != fraction] = 0\n\n    if class_reduction == \"micro\":\n        return fraction\n    if class_reduction == \"macro\":\n        return torch.mean(fraction)\n    if class_reduction == \"weighted\":\n        return torch.sum(fraction * (weights.float() / torch.sum(weights)))\n    if class_reduction == \"none\" or class_reduction is None:\n        return fraction\n\n    raise ValueError(f\"Reduction parameter {class_reduction} unknown. Choose between one of these: {valid_reduction}\")\n\n\ndef _simple_gather_all_tensors(result: Tensor, group: Any, world_size: int) -> List[Tensor]:\n    gathered_result = [torch.zeros_like(result) for _ in range(world_size)]\n    torch.distributed.all_gather(gathered_result, result, group)\n    return gathered_result\n\n\ndef gather_all_tensors(result: Tensor, group: Optional[Any] = None) -> List[Tensor]:\n    \"\"\"Gather all tensors from several ddp processes onto a list that is broadcasted to all processes.\n\n    Works on tensors that have the same number of dimensions, but where each dimension may differ. In this case\n    tensors are padded, gathered and then trimmed to secure equal workload for all processes.\n\n    Args:\n        result: the value to sync\n        group: the process group to gather results from. Defaults to all processes (world)\n\n    Return:\n        gathered_result: list with size equal to the process group where\n            ``gathered_result[i]`` corresponds to result tensor from process ``i``\n\n    \"\"\"\n    if group is None:\n        group = torch.distributed.group.WORLD\n\n    # convert tensors to contiguous format\n    result = result.contiguous()\n\n    world_size = torch.distributed.get_world_size(group)\n    torch.distributed.barrier(group=group)\n\n    # if the tensor is scalar, things are easy\n    if result.ndim == 0:\n        return _simple_gather_all_tensors(result, group, world_size)\n\n    # 1. Gather sizes of all tensors\n    local_size = torch.tensor(result.shape, device=result.device)\n    local_sizes = [torch.zeros_like(local_size) for _ in range(world_size)]\n    torch.distributed.all_gather(local_sizes, local_size, group=group)\n    max_size = torch.stack(local_sizes).max(dim=0).values\n    all_sizes_equal = all(all(ls == max_size) for ls in local_sizes)\n\n    # 2. If shapes are all the same, then do a simple gather:\n    if all_sizes_equal:\n        return _simple_gather_all_tensors(result, group, world_size)\n\n    # 3. If not, we need to pad each local tensor to maximum size, gather and then truncate\n    pad_dims = []\n    pad_by = (max_size - local_size).detach().cpu()\n    for val in reversed(pad_by):\n        pad_dims.append(0)\n        pad_dims.append(val.item())\n    result_padded = F.pad(result, pad_dims)\n    gathered_result = [torch.zeros_like(result_padded) for _ in range(world_size)]\n    torch.distributed.all_gather(gathered_result, result_padded, group)\n    for idx, item_size in enumerate(local_sizes):\n        slice_param = [slice(dim_size) for dim_size in item_size]\n        gathered_result[idx] = gathered_result[idx][slice_param]\n    return gathered_result\n", "path": "src/torchmetrics/utilities/distributed.py"}], "after_files": [{"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom torchmetrics.utilities.checks import check_forward_full_state_property\nfrom torchmetrics.utilities.data import (\n    dim_zero_cat,\n    dim_zero_max,\n    dim_zero_mean,\n    dim_zero_min,\n    dim_zero_sum,\n)\nfrom torchmetrics.utilities.distributed import class_reduce, reduce\nfrom torchmetrics.utilities.prints import rank_zero_debug, rank_zero_info, rank_zero_warn\n\n__all__ = [\n    \"check_forward_full_state_property\",\n    \"class_reduce\",\n    \"reduce\",\n    \"rank_zero_debug\",\n    \"rank_zero_info\",\n    \"rank_zero_warn\",\n    \"dim_zero_cat\",\n    \"dim_zero_max\",\n    \"dim_zero_mean\",\n    \"dim_zero_min\",\n    \"dim_zero_sum\",\n]\n", "path": "src/torchmetrics/utilities/__init__.py"}, {"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, List, Optional\n\nimport torch\nfrom torch import Tensor\nfrom torch.nn import functional as F  # noqa: N812\nfrom typing_extensions import Literal\n\n\ndef reduce(x: Tensor, reduction: Literal[\"elementwise_mean\", \"sum\", \"none\", None]) -> Tensor:\n    \"\"\"Reduces a given tensor by a given reduction method.\n\n    Args:\n        x: the tensor, which shall be reduced\n        reduction:  a string specifying the reduction method ('elementwise_mean', 'none', 'sum')\n\n    Return:\n        reduced Tensor\n\n    Raise:\n        ValueError if an invalid reduction parameter was given\n\n    \"\"\"\n    if reduction == \"elementwise_mean\":\n        return torch.mean(x)\n    if reduction == \"none\" or reduction is None:\n        return x\n    if reduction == \"sum\":\n        return torch.sum(x)\n    raise ValueError(\"Reduction parameter unknown.\")\n\n\ndef class_reduce(\n    num: Tensor,\n    denom: Tensor,\n    weights: Tensor,\n    class_reduction: Literal[\"micro\", \"macro\", \"weighted\", \"none\", None] = \"none\",\n) -> Tensor:\n    \"\"\"Reduce classification metrics of the form ``num / denom * weights``.\n\n    For example for calculating standard accuracy the num would be number of true positives per class, denom would be\n    the support per class, and weights would be a tensor of 1s.\n\n    Args:\n        num: numerator tensor\n        denom: denominator tensor\n        weights: weights for each class\n        class_reduction: reduction method for multiclass problems:\n\n            - ``'micro'``: calculate metrics globally (default)\n            - ``'macro'``: calculate metrics for each label, and find their unweighted mean.\n            - ``'weighted'``: calculate metrics for each label, and find their weighted mean.\n            - ``'none'`` or ``None``: returns calculated metric per class\n\n    Raises:\n        ValueError:\n            If ``class_reduction`` is none of ``\"micro\"``, ``\"macro\"``, ``\"weighted\"``, ``\"none\"`` or ``None``.\n\n    \"\"\"\n    valid_reduction = (\"micro\", \"macro\", \"weighted\", \"none\", None)\n    fraction = torch.sum(num) / torch.sum(denom) if class_reduction == \"micro\" else num / denom\n\n    # We need to take care of instances where the denom can be 0\n    # for some (or all) classes which will produce nans\n    fraction[fraction != fraction] = 0\n\n    if class_reduction == \"micro\":\n        return fraction\n    if class_reduction == \"macro\":\n        return torch.mean(fraction)\n    if class_reduction == \"weighted\":\n        return torch.sum(fraction * (weights.float() / torch.sum(weights)))\n    if class_reduction == \"none\" or class_reduction is None:\n        return fraction\n\n    raise ValueError(f\"Reduction parameter {class_reduction} unknown. Choose between one of these: {valid_reduction}\")\n\n\ndef _simple_gather_all_tensors(result: Tensor, group: Any, world_size: int) -> List[Tensor]:\n    gathered_result = [torch.zeros_like(result) for _ in range(world_size)]\n    torch.distributed.all_gather(gathered_result, result, group)\n    return gathered_result\n\n\ndef gather_all_tensors(result: Tensor, group: Optional[Any] = None) -> List[Tensor]:\n    \"\"\"Gather all tensors from several ddp processes onto a list that is broadcasted to all processes.\n\n    Works on tensors that have the same number of dimensions, but where each dimension may differ. In this case\n    tensors are padded, gathered and then trimmed to secure equal workload for all processes.\n\n    Args:\n        result: the value to sync\n        group: the process group to gather results from. Defaults to all processes (world)\n\n    Return:\n        list with size equal to the process group where element i corresponds to result tensor from process i\n\n    \"\"\"\n    if group is None:\n        group = torch.distributed.group.WORLD\n\n    # convert tensors to contiguous format\n    result = result.contiguous()\n\n    world_size = torch.distributed.get_world_size(group)\n    torch.distributed.barrier(group=group)\n\n    # if the tensor is scalar, things are easy\n    if result.ndim == 0:\n        return _simple_gather_all_tensors(result, group, world_size)\n\n    # 1. Gather sizes of all tensors\n    local_size = torch.tensor(result.shape, device=result.device)\n    local_sizes = [torch.zeros_like(local_size) for _ in range(world_size)]\n    torch.distributed.all_gather(local_sizes, local_size, group=group)\n    max_size = torch.stack(local_sizes).max(dim=0).values\n    all_sizes_equal = all(all(ls == max_size) for ls in local_sizes)\n\n    # 2. If shapes are all the same, then do a simple gather:\n    if all_sizes_equal:\n        return _simple_gather_all_tensors(result, group, world_size)\n\n    # 3. If not, we need to pad each local tensor to maximum size, gather and then truncate\n    pad_dims = []\n    pad_by = (max_size - local_size).detach().cpu()\n    for val in reversed(pad_by):\n        pad_dims.append(0)\n        pad_dims.append(val.item())\n    result_padded = F.pad(result, pad_dims)\n    gathered_result = [torch.zeros_like(result_padded) for _ in range(world_size)]\n    torch.distributed.all_gather(gathered_result, result_padded, group)\n    for idx, item_size in enumerate(local_sizes):\n        slice_param = [slice(dim_size) for dim_size in item_size]\n        gathered_result[idx] = gathered_result[idx][slice_param]\n    return gathered_result\n", "path": "src/torchmetrics/utilities/distributed.py"}]}