We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Minor discrepancy in AveragePrecision documentation for `average="none"`
## 📚 Documentation

The [documentation](https://torchmetrics.readthedocs.io/en/latest/classification/average_precision.html) for the `torchmetrics.AveragePrecision` class and `torchmetrics.functional.average_precision()` function state that setting `average="none"` is permitted. However, the source code only seems to allow `average=None` (see [here](https://github.com/Lightning-AI/metrics/blob/master/src/torchmetrics/classification/avg_precision.py#L98) and [here](https://github.com/Lightning-AI/metrics/blob/master/src/torchmetrics/functional/classification/average_precision.py#L175)).

I'd be happy to submit a PR for this but I would like some feedback on how to best resolve this, since I am not familiar with the design of this library. The two immediate directions I can think of are editing the documentation to only allow `average=None` or editing the source code to support `average="none"`.

Thanks!

--- 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/functional/classification/average_precision.py`
Content:
```
1 # Copyright The PyTorch 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 import warnings
15 from typing import List, Optional, Tuple, Union
16 
17 import torch
18 from torch import Tensor
19 
20 from torchmetrics.functional.classification.precision_recall_curve import (
21     _precision_recall_curve_compute,
22     _precision_recall_curve_update,
23 )
24 from torchmetrics.utilities.data import _bincount
25 
26 
27 def _average_precision_update(
28     preds: Tensor,
29     target: Tensor,
30     num_classes: Optional[int] = None,
31     pos_label: Optional[int] = None,
32     average: Optional[str] = "macro",
33 ) -> Tuple[Tensor, Tensor, int, Optional[int]]:
34     """Format the predictions and target based on the ``num_classes``, ``pos_label`` and ``average`` parameter.
35 
36     Args:
37         preds: predictions from model (logits or probabilities)
38         target: ground truth values
39         num_classes: integer with number of classes.
40         pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated
41             to 1. For multiclass problems this argument should not be set as we iteratively change it in the
42             range ``[0, num_classes-1]``
43         average: reduction method for multi-class or multi-label problems
44     """
45     preds, target, num_classes, pos_label = _precision_recall_curve_update(preds, target, num_classes, pos_label)
46     if average == "micro" and preds.ndim != target.ndim:
47         raise ValueError("Cannot use `micro` average with multi-class input")
48 
49     return preds, target, num_classes, pos_label
50 
51 
52 def _average_precision_compute(
53     preds: Tensor,
54     target: Tensor,
55     num_classes: int,
56     pos_label: Optional[int] = None,
57     average: Optional[str] = "macro",
58 ) -> Union[List[Tensor], Tensor]:
59     """Computes the average precision score.
60 
61     Args:
62         preds: predictions from model (logits or probabilities)
63         target: ground truth values
64         num_classes: integer with number of classes.
65         pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated
66             to 1. For multiclass problems his argument should not be set as we iteratively change it in the
67             range ``[0, num_classes-1]``
68         average: reduction method for multi-class or multi-label problems
69 
70     Example:
71         >>> # binary case
72         >>> preds = torch.tensor([0, 1, 2, 3])
73         >>> target = torch.tensor([0, 1, 1, 1])
74         >>> pos_label = 1
75         >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, pos_label=pos_label)
76         >>> _average_precision_compute(preds, target, num_classes, pos_label)
77         tensor(1.)
78 
79         >>> # multiclass case
80         >>> preds = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],
81         ...                      [0.05, 0.75, 0.05, 0.05, 0.05],
82         ...                      [0.05, 0.05, 0.75, 0.05, 0.05],
83         ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])
84         >>> target = torch.tensor([0, 1, 3, 2])
85         >>> num_classes = 5
86         >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes)
87         >>> _average_precision_compute(preds, target, num_classes, average=None)
88         [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]
89     """
90 
91     if average == "micro" and preds.ndim == target.ndim:
92         preds = preds.flatten()
93         target = target.flatten()
94         num_classes = 1
95 
96     precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes, pos_label)
97     if average == "weighted":
98         if preds.ndim == target.ndim and target.ndim > 1:
99             weights = target.sum(dim=0).float()
100         else:
101             weights = _bincount(target, minlength=num_classes).float()
102         weights = weights / torch.sum(weights)
103     else:
104         weights = None
105     return _average_precision_compute_with_precision_recall(precision, recall, num_classes, average, weights)
106 
107 
108 def _average_precision_compute_with_precision_recall(
109     precision: Tensor,
110     recall: Tensor,
111     num_classes: int,
112     average: Optional[str] = "macro",
113     weights: Optional[Tensor] = None,
114 ) -> Union[List[Tensor], Tensor]:
115     """Computes the average precision score from precision and recall.
116 
117     Args:
118         precision: precision values
119         recall: recall values
120         num_classes: integer with number of classes. Not nessesary to provide
121             for binary problems.
122         average: reduction method for multi-class or multi-label problems
123         weights: weights to use when average='weighted'
124 
125     Example:
126         >>> # binary case
127         >>> preds = torch.tensor([0, 1, 2, 3])
128         >>> target = torch.tensor([0, 1, 1, 1])
129         >>> pos_label = 1
130         >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, pos_label=pos_label)
131         >>> precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes, pos_label)
132         >>> _average_precision_compute_with_precision_recall(precision, recall, num_classes, average=None)
133         tensor(1.)
134 
135         >>> # multiclass case
136         >>> preds = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],
137         ...                      [0.05, 0.75, 0.05, 0.05, 0.05],
138         ...                      [0.05, 0.05, 0.75, 0.05, 0.05],
139         ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])
140         >>> target = torch.tensor([0, 1, 3, 2])
141         >>> num_classes = 5
142         >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes)
143         >>> precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes)
144         >>> _average_precision_compute_with_precision_recall(precision, recall, num_classes, average=None)
145         [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]
146     """
147 
148     # Return the step function integral
149     # The following works because the last entry of precision is
150     # guaranteed to be 1, as returned by precision_recall_curve
151     if num_classes == 1:
152         return -torch.sum((recall[1:] - recall[:-1]) * precision[:-1])
153 
154     res = []
155     for p, r in zip(precision, recall):
156         res.append(-torch.sum((r[1:] - r[:-1]) * p[:-1]))
157 
158     # Reduce
159     if average in ("macro", "weighted"):
160         res = torch.stack(res)
161         if torch.isnan(res).any():
162             warnings.warn(
163                 "Average precision score for one or more classes was `nan`. Ignoring these classes in average",
164                 UserWarning,
165             )
166         if average == "macro":
167             return res[~torch.isnan(res)].mean()
168         weights = torch.ones_like(res) if weights is None else weights
169         return (res * weights)[~torch.isnan(res)].sum()
170     if average is None:
171         return res
172     allowed_average = ("micro", "macro", "weighted", None)
173     raise ValueError(f"Expected argument `average` to be one of {allowed_average}" f" but got {average}")
174 
175 
176 def average_precision(
177     preds: Tensor,
178     target: Tensor,
179     num_classes: Optional[int] = None,
180     pos_label: Optional[int] = None,
181     average: Optional[str] = "macro",
182 ) -> Union[List[Tensor], Tensor]:
183     """Computes the average precision score.
184 
185     Args:
186         preds: predictions from model (logits or probabilities)
187         target: ground truth values
188         num_classes: integer with number of classes. Not nessesary to provide
189             for binary problems.
190         pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated
191             to 1. For multiclass problems his argument should not be set as we iteratively change it in the
192             range ``[0, num_classes-1]``
193         average:
194             defines the reduction that is applied in the case of multiclass and multilabel input.
195             Should be one of the following:
196 
197             - ``'macro'`` [default]: Calculate the metric for each class separately, and average the
198               metrics across classes (with equal weights for each class).
199             - ``'micro'``: Calculate the metric globally, across all samples and classes. Cannot be
200               used with multiclass input.
201             - ``'weighted'``: Calculate the metric for each class separately, and average the
202               metrics across classes, weighting each class by its support.
203             - ``'none'`` or ``None``: Calculate the metric for each class separately, and return
204               the metric for every class.
205 
206     Returns:
207         tensor with average precision. If multiclass will return list
208         of such tensors, one for each class
209 
210     Example (binary case):
211         >>> from torchmetrics.functional import average_precision
212         >>> pred = torch.tensor([0, 1, 2, 3])
213         >>> target = torch.tensor([0, 1, 1, 1])
214         >>> average_precision(pred, target, pos_label=1)
215         tensor(1.)
216 
217     Example (multiclass case):
218         >>> pred = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],
219         ...                      [0.05, 0.75, 0.05, 0.05, 0.05],
220         ...                      [0.05, 0.05, 0.75, 0.05, 0.05],
221         ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])
222         >>> target = torch.tensor([0, 1, 3, 2])
223         >>> average_precision(pred, target, num_classes=5, average=None)
224         [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]
225     """
226     preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes, pos_label, average)
227     return _average_precision_compute(preds, target, num_classes, pos_label, average)
228
```
--- 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/functional/classification/average_precision.py b/src/torchmetrics/functional/classification/average_precision.py
--- a/src/torchmetrics/functional/classification/average_precision.py
+++ b/src/torchmetrics/functional/classification/average_precision.py
@@ -167,9 +167,9 @@
             return res[~torch.isnan(res)].mean()
         weights = torch.ones_like(res) if weights is None else weights
         return (res * weights)[~torch.isnan(res)].sum()
-    if average is None:
+    if average is None or average == "none":
         return res
-    allowed_average = ("micro", "macro", "weighted", None)
+    allowed_average = ("micro", "macro", "weighted", "none", None)
     raise ValueError(f"Expected argument `average` to be one of {allowed_average}" f" but got {average}")

{"golden_diff": "diff --git a/src/torchmetrics/functional/classification/average_precision.py b/src/torchmetrics/functional/classification/average_precision.py\n--- a/src/torchmetrics/functional/classification/average_precision.py\n+++ b/src/torchmetrics/functional/classification/average_precision.py\n@@ -167,9 +167,9 @@\n             return res[~torch.isnan(res)].mean()\n         weights = torch.ones_like(res) if weights is None else weights\n         return (res * weights)[~torch.isnan(res)].sum()\n-    if average is None:\n+    if average is None or average == \"none\":\n         return res\n-    allowed_average = (\"micro\", \"macro\", \"weighted\", None)\n+    allowed_average = (\"micro\", \"macro\", \"weighted\", \"none\", None)\n     raise ValueError(f\"Expected argument `average` to be one of {allowed_average}\" f\" but got {average}\")\n", "issue": "Minor discrepancy in AveragePrecision documentation for `average=\"none\"`\n## \ud83d\udcda Documentation\r\n\r\nThe [documentation](https://torchmetrics.readthedocs.io/en/latest/classification/average_precision.html) for the `torchmetrics.AveragePrecision` class and `torchmetrics.functional.average_precision()` function state that setting `average=\"none\"` is permitted. However, the source code only seems to allow `average=None` (see [here](https://github.com/Lightning-AI/metrics/blob/master/src/torchmetrics/classification/avg_precision.py#L98) and [here](https://github.com/Lightning-AI/metrics/blob/master/src/torchmetrics/functional/classification/average_precision.py#L175)).\r\n\r\nI'd be happy to submit a PR for this but I would like some feedback on how to best resolve this, since I am not familiar with the design of this library. The two immediate directions I can think of are editing the documentation to only allow `average=None` or editing the source code to support `average=\"none\"`.\r\n\r\nThanks!\n", "before_files": [{"content": "# Copyright The PyTorch 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.\nimport warnings\nfrom typing import List, Optional, Tuple, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torchmetrics.functional.classification.precision_recall_curve import (\n    _precision_recall_curve_compute,\n    _precision_recall_curve_update,\n)\nfrom torchmetrics.utilities.data import _bincount\n\n\ndef _average_precision_update(\n    preds: Tensor,\n    target: Tensor,\n    num_classes: Optional[int] = None,\n    pos_label: Optional[int] = None,\n    average: Optional[str] = \"macro\",\n) -> Tuple[Tensor, Tensor, int, Optional[int]]:\n    \"\"\"Format the predictions and target based on the ``num_classes``, ``pos_label`` and ``average`` parameter.\n\n    Args:\n        preds: predictions from model (logits or probabilities)\n        target: ground truth values\n        num_classes: integer with number of classes.\n        pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated\n            to 1. For multiclass problems this argument should not be set as we iteratively change it in the\n            range ``[0, num_classes-1]``\n        average: reduction method for multi-class or multi-label problems\n    \"\"\"\n    preds, target, num_classes, pos_label = _precision_recall_curve_update(preds, target, num_classes, pos_label)\n    if average == \"micro\" and preds.ndim != target.ndim:\n        raise ValueError(\"Cannot use `micro` average with multi-class input\")\n\n    return preds, target, num_classes, pos_label\n\n\ndef _average_precision_compute(\n    preds: Tensor,\n    target: Tensor,\n    num_classes: int,\n    pos_label: Optional[int] = None,\n    average: Optional[str] = \"macro\",\n) -> Union[List[Tensor], Tensor]:\n    \"\"\"Computes the average precision score.\n\n    Args:\n        preds: predictions from model (logits or probabilities)\n        target: ground truth values\n        num_classes: integer with number of classes.\n        pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated\n            to 1. For multiclass problems his argument should not be set as we iteratively change it in the\n            range ``[0, num_classes-1]``\n        average: reduction method for multi-class or multi-label problems\n\n    Example:\n        >>> # binary case\n        >>> preds = torch.tensor([0, 1, 2, 3])\n        >>> target = torch.tensor([0, 1, 1, 1])\n        >>> pos_label = 1\n        >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, pos_label=pos_label)\n        >>> _average_precision_compute(preds, target, num_classes, pos_label)\n        tensor(1.)\n\n        >>> # multiclass case\n        >>> preds = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])\n        >>> target = torch.tensor([0, 1, 3, 2])\n        >>> num_classes = 5\n        >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes)\n        >>> _average_precision_compute(preds, target, num_classes, average=None)\n        [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]\n    \"\"\"\n\n    if average == \"micro\" and preds.ndim == target.ndim:\n        preds = preds.flatten()\n        target = target.flatten()\n        num_classes = 1\n\n    precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes, pos_label)\n    if average == \"weighted\":\n        if preds.ndim == target.ndim and target.ndim > 1:\n            weights = target.sum(dim=0).float()\n        else:\n            weights = _bincount(target, minlength=num_classes).float()\n        weights = weights / torch.sum(weights)\n    else:\n        weights = None\n    return _average_precision_compute_with_precision_recall(precision, recall, num_classes, average, weights)\n\n\ndef _average_precision_compute_with_precision_recall(\n    precision: Tensor,\n    recall: Tensor,\n    num_classes: int,\n    average: Optional[str] = \"macro\",\n    weights: Optional[Tensor] = None,\n) -> Union[List[Tensor], Tensor]:\n    \"\"\"Computes the average precision score from precision and recall.\n\n    Args:\n        precision: precision values\n        recall: recall values\n        num_classes: integer with number of classes. Not nessesary to provide\n            for binary problems.\n        average: reduction method for multi-class or multi-label problems\n        weights: weights to use when average='weighted'\n\n    Example:\n        >>> # binary case\n        >>> preds = torch.tensor([0, 1, 2, 3])\n        >>> target = torch.tensor([0, 1, 1, 1])\n        >>> pos_label = 1\n        >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, pos_label=pos_label)\n        >>> precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes, pos_label)\n        >>> _average_precision_compute_with_precision_recall(precision, recall, num_classes, average=None)\n        tensor(1.)\n\n        >>> # multiclass case\n        >>> preds = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])\n        >>> target = torch.tensor([0, 1, 3, 2])\n        >>> num_classes = 5\n        >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes)\n        >>> precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes)\n        >>> _average_precision_compute_with_precision_recall(precision, recall, num_classes, average=None)\n        [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]\n    \"\"\"\n\n    # Return the step function integral\n    # The following works because the last entry of precision is\n    # guaranteed to be 1, as returned by precision_recall_curve\n    if num_classes == 1:\n        return -torch.sum((recall[1:] - recall[:-1]) * precision[:-1])\n\n    res = []\n    for p, r in zip(precision, recall):\n        res.append(-torch.sum((r[1:] - r[:-1]) * p[:-1]))\n\n    # Reduce\n    if average in (\"macro\", \"weighted\"):\n        res = torch.stack(res)\n        if torch.isnan(res).any():\n            warnings.warn(\n                \"Average precision score for one or more classes was `nan`. Ignoring these classes in average\",\n                UserWarning,\n            )\n        if average == \"macro\":\n            return res[~torch.isnan(res)].mean()\n        weights = torch.ones_like(res) if weights is None else weights\n        return (res * weights)[~torch.isnan(res)].sum()\n    if average is None:\n        return res\n    allowed_average = (\"micro\", \"macro\", \"weighted\", None)\n    raise ValueError(f\"Expected argument `average` to be one of {allowed_average}\" f\" but got {average}\")\n\n\ndef average_precision(\n    preds: Tensor,\n    target: Tensor,\n    num_classes: Optional[int] = None,\n    pos_label: Optional[int] = None,\n    average: Optional[str] = \"macro\",\n) -> Union[List[Tensor], Tensor]:\n    \"\"\"Computes the average precision score.\n\n    Args:\n        preds: predictions from model (logits or probabilities)\n        target: ground truth values\n        num_classes: integer with number of classes. Not nessesary to provide\n            for binary problems.\n        pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated\n            to 1. For multiclass problems his argument should not be set as we iteratively change it in the\n            range ``[0, num_classes-1]``\n        average:\n            defines the reduction that is applied in the case of multiclass and multilabel input.\n            Should be one of the following:\n\n            - ``'macro'`` [default]: Calculate the metric for each class separately, and average the\n              metrics across classes (with equal weights for each class).\n            - ``'micro'``: Calculate the metric globally, across all samples and classes. Cannot be\n              used with multiclass input.\n            - ``'weighted'``: Calculate the metric for each class separately, and average the\n              metrics across classes, weighting each class by its support.\n            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return\n              the metric for every class.\n\n    Returns:\n        tensor with average precision. If multiclass will return list\n        of such tensors, one for each class\n\n    Example (binary case):\n        >>> from torchmetrics.functional import average_precision\n        >>> pred = torch.tensor([0, 1, 2, 3])\n        >>> target = torch.tensor([0, 1, 1, 1])\n        >>> average_precision(pred, target, pos_label=1)\n        tensor(1.)\n\n    Example (multiclass case):\n        >>> pred = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])\n        >>> target = torch.tensor([0, 1, 3, 2])\n        >>> average_precision(pred, target, num_classes=5, average=None)\n        [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]\n    \"\"\"\n    preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes, pos_label, average)\n    return _average_precision_compute(preds, target, num_classes, pos_label, average)\n", "path": "src/torchmetrics/functional/classification/average_precision.py"}], "after_files": [{"content": "# Copyright The PyTorch 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.\nimport warnings\nfrom typing import List, Optional, Tuple, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torchmetrics.functional.classification.precision_recall_curve import (\n    _precision_recall_curve_compute,\n    _precision_recall_curve_update,\n)\nfrom torchmetrics.utilities.data import _bincount\n\n\ndef _average_precision_update(\n    preds: Tensor,\n    target: Tensor,\n    num_classes: Optional[int] = None,\n    pos_label: Optional[int] = None,\n    average: Optional[str] = \"macro\",\n) -> Tuple[Tensor, Tensor, int, Optional[int]]:\n    \"\"\"Format the predictions and target based on the ``num_classes``, ``pos_label`` and ``average`` parameter.\n\n    Args:\n        preds: predictions from model (logits or probabilities)\n        target: ground truth values\n        num_classes: integer with number of classes.\n        pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated\n            to 1. For multiclass problems this argument should not be set as we iteratively change it in the\n            range ``[0, num_classes-1]``\n        average: reduction method for multi-class or multi-label problems\n    \"\"\"\n    preds, target, num_classes, pos_label = _precision_recall_curve_update(preds, target, num_classes, pos_label)\n    if average == \"micro\" and preds.ndim != target.ndim:\n        raise ValueError(\"Cannot use `micro` average with multi-class input\")\n\n    return preds, target, num_classes, pos_label\n\n\ndef _average_precision_compute(\n    preds: Tensor,\n    target: Tensor,\n    num_classes: int,\n    pos_label: Optional[int] = None,\n    average: Optional[str] = \"macro\",\n) -> Union[List[Tensor], Tensor]:\n    \"\"\"Computes the average precision score.\n\n    Args:\n        preds: predictions from model (logits or probabilities)\n        target: ground truth values\n        num_classes: integer with number of classes.\n        pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated\n            to 1. For multiclass problems his argument should not be set as we iteratively change it in the\n            range ``[0, num_classes-1]``\n        average: reduction method for multi-class or multi-label problems\n\n    Example:\n        >>> # binary case\n        >>> preds = torch.tensor([0, 1, 2, 3])\n        >>> target = torch.tensor([0, 1, 1, 1])\n        >>> pos_label = 1\n        >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, pos_label=pos_label)\n        >>> _average_precision_compute(preds, target, num_classes, pos_label)\n        tensor(1.)\n\n        >>> # multiclass case\n        >>> preds = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])\n        >>> target = torch.tensor([0, 1, 3, 2])\n        >>> num_classes = 5\n        >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes)\n        >>> _average_precision_compute(preds, target, num_classes, average=None)\n        [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]\n    \"\"\"\n\n    if average == \"micro\" and preds.ndim == target.ndim:\n        preds = preds.flatten()\n        target = target.flatten()\n        num_classes = 1\n\n    precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes, pos_label)\n    if average == \"weighted\":\n        if preds.ndim == target.ndim and target.ndim > 1:\n            weights = target.sum(dim=0).float()\n        else:\n            weights = _bincount(target, minlength=num_classes).float()\n        weights = weights / torch.sum(weights)\n    else:\n        weights = None\n    return _average_precision_compute_with_precision_recall(precision, recall, num_classes, average, weights)\n\n\ndef _average_precision_compute_with_precision_recall(\n    precision: Tensor,\n    recall: Tensor,\n    num_classes: int,\n    average: Optional[str] = \"macro\",\n    weights: Optional[Tensor] = None,\n) -> Union[List[Tensor], Tensor]:\n    \"\"\"Computes the average precision score from precision and recall.\n\n    Args:\n        precision: precision values\n        recall: recall values\n        num_classes: integer with number of classes. Not nessesary to provide\n            for binary problems.\n        average: reduction method for multi-class or multi-label problems\n        weights: weights to use when average='weighted'\n\n    Example:\n        >>> # binary case\n        >>> preds = torch.tensor([0, 1, 2, 3])\n        >>> target = torch.tensor([0, 1, 1, 1])\n        >>> pos_label = 1\n        >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, pos_label=pos_label)\n        >>> precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes, pos_label)\n        >>> _average_precision_compute_with_precision_recall(precision, recall, num_classes, average=None)\n        tensor(1.)\n\n        >>> # multiclass case\n        >>> preds = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])\n        >>> target = torch.tensor([0, 1, 3, 2])\n        >>> num_classes = 5\n        >>> preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes)\n        >>> precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes)\n        >>> _average_precision_compute_with_precision_recall(precision, recall, num_classes, average=None)\n        [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]\n    \"\"\"\n\n    # Return the step function integral\n    # The following works because the last entry of precision is\n    # guaranteed to be 1, as returned by precision_recall_curve\n    if num_classes == 1:\n        return -torch.sum((recall[1:] - recall[:-1]) * precision[:-1])\n\n    res = []\n    for p, r in zip(precision, recall):\n        res.append(-torch.sum((r[1:] - r[:-1]) * p[:-1]))\n\n    # Reduce\n    if average in (\"macro\", \"weighted\"):\n        res = torch.stack(res)\n        if torch.isnan(res).any():\n            warnings.warn(\n                \"Average precision score for one or more classes was `nan`. Ignoring these classes in average\",\n                UserWarning,\n            )\n        if average == \"macro\":\n            return res[~torch.isnan(res)].mean()\n        weights = torch.ones_like(res) if weights is None else weights\n        return (res * weights)[~torch.isnan(res)].sum()\n    if average is None or average == \"none\":\n        return res\n    allowed_average = (\"micro\", \"macro\", \"weighted\", \"none\", None)\n    raise ValueError(f\"Expected argument `average` to be one of {allowed_average}\" f\" but got {average}\")\n\n\ndef average_precision(\n    preds: Tensor,\n    target: Tensor,\n    num_classes: Optional[int] = None,\n    pos_label: Optional[int] = None,\n    average: Optional[str] = \"macro\",\n) -> Union[List[Tensor], Tensor]:\n    \"\"\"Computes the average precision score.\n\n    Args:\n        preds: predictions from model (logits or probabilities)\n        target: ground truth values\n        num_classes: integer with number of classes. Not nessesary to provide\n            for binary problems.\n        pos_label: integer determining the positive class. Default is ``None`` which for binary problem is translated\n            to 1. For multiclass problems his argument should not be set as we iteratively change it in the\n            range ``[0, num_classes-1]``\n        average:\n            defines the reduction that is applied in the case of multiclass and multilabel input.\n            Should be one of the following:\n\n            - ``'macro'`` [default]: Calculate the metric for each class separately, and average the\n              metrics across classes (with equal weights for each class).\n            - ``'micro'``: Calculate the metric globally, across all samples and classes. Cannot be\n              used with multiclass input.\n            - ``'weighted'``: Calculate the metric for each class separately, and average the\n              metrics across classes, weighting each class by its support.\n            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return\n              the metric for every class.\n\n    Returns:\n        tensor with average precision. If multiclass will return list\n        of such tensors, one for each class\n\n    Example (binary case):\n        >>> from torchmetrics.functional import average_precision\n        >>> pred = torch.tensor([0, 1, 2, 3])\n        >>> target = torch.tensor([0, 1, 1, 1])\n        >>> average_precision(pred, target, pos_label=1)\n        tensor(1.)\n\n    Example (multiclass case):\n        >>> pred = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],\n        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])\n        >>> target = torch.tensor([0, 1, 3, 2])\n        >>> average_precision(pred, target, num_classes=5, average=None)\n        [tensor(1.), tensor(1.), tensor(0.2500), tensor(0.2500), tensor(nan)]\n    \"\"\"\n    preds, target, num_classes, pos_label = _average_precision_update(preds, target, num_classes, pos_label, average)\n    return _average_precision_compute(preds, target, num_classes, pos_label, average)\n", "path": "src/torchmetrics/functional/classification/average_precision.py"}]}