Repository: Lightning-AI/torchmetrics

Issue:
AUC computation is incorrect for tensors with >2048 elements.
## 🐛 Bug

When `auc` is called with `reorder=True`, a stable sort is performed on the input tensors using `_stable_1d_sort` ([code](https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/functional/classification/auc.py#L38)). However, `_stable_1d_sort` returns at most `2048` elements if `nb` is not provided ([code](https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/utilities/data.py#L179)). This leads to `auc` computation being incorrect. 
Make `_stable_1d_sort(nb)` optional
# Before submitting

- [x] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), Pull Request section?
- [n/a] Did you make sure to update the docs?   
- [x] Did you write any new necessary tests?  

## What does this PR do?
Fixes #195


Discussion:
Hi! thanks for your contribution!, great first issue!
# [Codecov](https://codecov.io/gh/PyTorchLightning/metrics/pull/196?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) Report
> Merging [#196](https://codecov.io/gh/PyTorchLightning/metrics/pull/196?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) (25c8a2d) into [master](https://codecov.io/gh/PyTorchLightning/metrics/commit/d3058c4e60f3681c74fa5b28a56a52fcabe9698d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) (d3058c4) will **decrease** coverage by `16.76%`.
> The diff coverage is `75.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/graphs/tree.svg?width=650&height=150&src=pr&token=NER6LPI3HS&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning)](https://codecov.io/gh/PyTorchLightning/metrics/pull/196?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning)

```diff
@@             Coverage Diff             @@
##           master     #196       +/-   ##
===========================================
- Coverage   96.63%   79.86%   -16.77%     
===========================================
  Files         180       90       -90     
  Lines        5760     2871     -2889     
===========================================
- Hits         5566     2293     -3273     
- Misses        194      578      +384     
```

| Flag | Coverage Δ | |
|---|---|---|
| Linux | `79.86% <75.00%> (-0.07%)` | :arrow_down: |
| Windows | `79.86% <75.00%> (-0.07%)` | :arrow_down: |
| cpu | `79.86% <75.00%> (-16.77%)` | :arrow_down: |
| gpu | `?` | |
| macOS | `79.86% <75.00%> (-16.77%)` | :arrow_down: |
| pytest | `79.86% <75.00%> (-16.77%)` | :arrow_down: |
| python3.6 | `?` | |
| python3.8 | `?` | |
| python3.9 | `?` | |
| torch1.3.1 | `?` | |
| torch1.4.0 | `?` | |
| torch1.8.1 | `?` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/PyTorchLightning/metrics/pull/196?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) | Coverage Δ | |
|---|---|---|
| [torchmetrics/utilities/data.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL3V0aWxpdGllcy9kYXRhLnB5) | `82.08% <75.00%> (-13.37%)` | :arrow_down: |
| [torchmetrics/utilities/distributed.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL3V0aWxpdGllcy9kaXN0cmlidXRlZC5weQ==) | `22.85% <0.00%> (-74.29%)` | :arrow_down: |
| [torchmetrics/classification/auc.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL2NsYXNzaWZpY2F0aW9uL2F1Yy5weQ==) | `47.61% <0.00%> (-52.39%)` | :arrow_down: |
| [torchmetrics/functional/classification/auroc.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL2Z1bmN0aW9uYWwvY2xhc3NpZmljYXRpb24vYXVyb2MucHk=) | `46.15% <0.00%> (-40.01%)` | :arrow_down: |
| [torchmetrics/metric.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL21ldHJpYy5weQ==) | `56.17% <0.00%> (-39.05%)` | :arrow_down: |
| [torchmetrics/functional/regression/psnr.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL2Z1bmN0aW9uYWwvcmVncmVzc2lvbi9wc25yLnB5) | `60.60% <0.00%> (-36.37%)` | :arrow_down: |
| [torchmetrics/functional/classification/accuracy.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL2Z1bmN0aW9uYWwvY2xhc3NpZmljYXRpb24vYWNjdXJhY3kucHk=) | `57.57% <0.00%> (-36.37%)` | :arrow_down: |
| [...chmetrics/functional/classification/stat\_scores.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL2Z1bmN0aW9uYWwvY2xhc3NpZmljYXRpb24vc3RhdF9zY29yZXMucHk=) | `66.66% <0.00%> (-33.34%)` | :arrow_down: |
| [...rics/functional/classification/confusion\_matrix.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL2Z1bmN0aW9uYWwvY2xhc3NpZmljYXRpb24vY29uZnVzaW9uX21hdHJpeC5weQ==) | `66.66% <0.00%> (-33.34%)` | :arrow_down: |
| [torchmetrics/utilities/checks.py](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-dG9yY2htZXRyaWNzL3V0aWxpdGllcy9jaGVja3MucHk=) | `60.54% <0.00%> (-31.90%)` | :arrow_down: |
| ... and [133 more](https://codecov.io/gh/PyTorchLightning/metrics/pull/196/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PyTorchLightning/metrics/pull/196?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PyTorchLightning/metrics/pull/196?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning). Last update [d3058c4...25c8a2d](https://codecov.io/gh/PyTorchLightning/metrics/pull/196?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning).


Resolving patch:
diff --git a/torchmetrics/functional/classification/auc.py b/torchmetrics/functional/classification/auc.py
--- a/torchmetrics/functional/classification/auc.py
+++ b/torchmetrics/functional/classification/auc.py
@@ -16,8 +16,6 @@
 import torch
 from torch import Tensor
 
-from torchmetrics.utilities.data import _stable_1d_sort
-
 
 def _auc_update(x: Tensor, y: Tensor) -> Tuple[Tensor, Tensor]:
     if x.ndim > 1 or y.ndim > 1:
@@ -35,7 +33,8 @@ def _auc_update(x: Tensor, y: Tensor) -> Tuple[Tensor, Tensor]:
 
 def _auc_compute(x: Tensor, y: Tensor, reorder: bool = False) -> Tensor:
     if reorder:
-        x, x_idx = _stable_1d_sort(x)
+        # TODO: include stable=True arg when pytorch v1.9 is released
+        x, x_idx = torch.sort(x)
         y = y[x_idx]
 
     dx = x[1:] - x[:-1]
diff --git a/torchmetrics/utilities/data.py b/torchmetrics/utilities/data.py
--- a/torchmetrics/utilities/data.py
+++ b/torchmetrics/utilities/data.py
@@ -151,35 +151,6 @@ def get_num_classes(
     return num_classes
 
 
-def _stable_1d_sort(x: torch, nb: int = 2049):
-    """
-    Stable sort of 1d tensors. Pytorch defaults to a stable sorting algorithm
-    if number of elements are larger than 2048. This function pads the tensors,
-    makes the sort and returns the sorted array (with the padding removed)
-    See this discussion: https://discuss.pytorch.org/t/is-torch-sort-stable/20714
-
-    Raises:
-        ValueError:
-            If dim of ``x`` is greater than 1 since stable sort works with only 1d tensors.
-
-    Example:
-        >>> data = torch.tensor([8, 7, 2, 6, 4, 5, 3, 1, 9, 0])
-        >>> _stable_1d_sort(data)
-        (tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), tensor([9, 7, 2, 6, 4, 5, 3, 1, 0, 8]))
-        >>> _stable_1d_sort(data, nb=5)
-        (tensor([0, 1, 2, 3, 4]), tensor([9, 7, 2, 6, 4]))
-    """
-    if x.ndim > 1:
-        raise ValueError('Stable sort only works on 1d tensors')
-    n = x.numel()
-    if n < nb:
-        x_max = x.max()
-        x = torch.cat([x, (x_max + 1) * torch.ones(nb - n, dtype=x.dtype, device=x.device)], 0)
-    x_sort = x.sort()
-    i = min(nb, n)
-    return x_sort.values[:i], x_sort.indices[:i]
-
-
 def apply_to_collection(
     data: Any,
     dtype: Union[type, tuple],


Test patch:
diff --git a/tests/classification/test_auc.py b/tests/classification/test_auc.py
--- a/tests/classification/test_auc.py
+++ b/tests/classification/test_auc.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections import namedtuple
+from functools import partial
 
 import numpy as np
 import pytest
@@ -26,9 +27,13 @@
 seed_all(42)
 
 
-def sk_auc(x, y):
+def sk_auc(x, y, reorder=False):
     x = x.flatten()
     y = y.flatten()
+    if reorder:
+        idx = np.argsort(x, kind='stable')
+        x = x[idx]
+        y = y[idx]
     return _sk_auc(x, y)
 
 
@@ -36,15 +41,16 @@ def sk_auc(x, y):
 
 _examples = []
 # generate already ordered samples, sorted in both directions
-for i in range(4):
-    x = np.random.randint(0, 5, (NUM_BATCHES * 8))
-    y = np.random.randint(0, 5, (NUM_BATCHES * 8))
-    idx = np.argsort(x, kind='stable')
-    x = x[idx] if i % 2 == 0 else x[idx[::-1]]
-    y = y[idx] if i % 2 == 0 else x[idx[::-1]]
-    x = x.reshape(NUM_BATCHES, 8)
-    y = y.reshape(NUM_BATCHES, 8)
-    _examples.append(Input(x=tensor(x), y=tensor(y)))
+for batch_size in (8, 4049):
+    for i in range(4):
+        x = np.random.rand((NUM_BATCHES * batch_size))
+        y = np.random.rand((NUM_BATCHES * batch_size))
+        idx = np.argsort(x, kind='stable')
+        x = x[idx] if i % 2 == 0 else x[idx[::-1]]
+        y = y[idx] if i % 2 == 0 else x[idx[::-1]]
+        x = x.reshape(NUM_BATCHES, batch_size)
+        y = y.reshape(NUM_BATCHES, batch_size)
+        _examples.append(Input(x=tensor(x), y=tensor(y)))
 
 
 @pytest.mark.parametrize("x, y", _examples)
@@ -62,8 +68,11 @@ def test_auc(self, x, y, ddp, dist_sync_on_step):
             dist_sync_on_step=dist_sync_on_step,
         )
 
-    def test_auc_functional(self, x, y):
-        self.run_functional_metric_test(x, y, metric_functional=auc, sk_metric=sk_auc, metric_args={"reorder": False})
+    @pytest.mark.parametrize("reorder", [True, False])
+    def test_auc_functional(self, x, y, reorder):
+        self.run_functional_metric_test(
+            x, y, metric_functional=auc, sk_metric=partial(sk_auc, reorder=reorder), metric_args={"reorder": reorder}
+        )
 
 
 @pytest.mark.parametrize(['x', 'y', 'expected'], [