Repository: Lightning-AI/torchmetrics

Issue:
Half precision doesn't work with top_k > 1
## 🐛 Bug

When using predictions at half-precision, MulticlassAccuracy throws an error when using top_k > 1. Please see the code below:

### To Reproduce

Steps to reproduce the behavior...

```py
import torch
from torchmetrics.classification import MulticlassAccuracy

preds = torch.tensor([[1, 0, 0]]).half()
targets = torch.tensor([1])

metric = MulticlassAccuracy(num_classes=3, top_k=2)
metric(preds, targets)
metric.compute()
```


### Environment

- TorchMetrics version: 1.0.1
- Python & PyTorch Version (e.g., 1.0): 3.10 and 2.0.1




Discussion:
Most likely related issue on pytorch side:
https://github.com/pytorch/pytorch/issues/58841

Resolving patch:
diff --git a/src/torchmetrics/utilities/data.py b/src/torchmetrics/utilities/data.py
--- a/src/torchmetrics/utilities/data.py
+++ b/src/torchmetrics/utilities/data.py
@@ -112,6 +112,14 @@ def to_onehot(
     return tensor_onehot.scatter_(1, index, 1.0)
 
 
+def _top_k_with_half_precision_support(x: Tensor, k: int = 1, dim: int = 1) -> Tensor:
+    """torch.top_k does not support half precision on CPU."""
+    if x.dtype == torch.half and not x.is_cuda:
+        idx = torch.argsort(x, dim=dim, descending=True)
+        return idx.narrow(dim, 0, k)
+    return x.topk(k=k, dim=dim).indices
+
+
 def select_topk(prob_tensor: Tensor, topk: int = 1, dim: int = 1) -> Tensor:
     """Convert a probability tensor to binary by selecting top-k the highest entries.
 
@@ -135,7 +143,7 @@ def select_topk(prob_tensor: Tensor, topk: int = 1, dim: int = 1) -> Tensor:
     if topk == 1:  # argmax has better performance than topk
         topk_tensor = zeros.scatter(dim, prob_tensor.argmax(dim=dim, keepdim=True), 1.0)
     else:
-        topk_tensor = zeros.scatter(dim, prob_tensor.topk(k=topk, dim=dim).indices, 1.0)
+        topk_tensor = zeros.scatter(dim, _top_k_with_half_precision_support(prob_tensor, k=topk, dim=dim), 1.0)
     return topk_tensor.int()
 
 


Test patch:
diff --git a/tests/unittests/utilities/test_utilities.py b/tests/unittests/utilities/test_utilities.py
--- a/tests/unittests/utilities/test_utilities.py
+++ b/tests/unittests/utilities/test_utilities.py
@@ -20,7 +20,15 @@
 from torchmetrics.regression import MeanSquaredError, PearsonCorrCoef
 from torchmetrics.utilities import check_forward_full_state_property, rank_zero_debug, rank_zero_info, rank_zero_warn
 from torchmetrics.utilities.checks import _allclose_recursive
-from torchmetrics.utilities.data import _bincount, _cumsum, _flatten, _flatten_dict, to_categorical, to_onehot
+from torchmetrics.utilities.data import (
+    _bincount,
+    _cumsum,
+    _flatten,
+    _flatten_dict,
+    select_topk,
+    to_categorical,
+    to_onehot,
+)
 from torchmetrics.utilities.distributed import class_reduce, reduce
 from torchmetrics.utilities.exceptions import TorchMetricsUserWarning
 from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_13
@@ -201,3 +209,38 @@ def test_custom_cumsum():
         res = _cumsum(x, dim=0).cpu()
     res2 = np.cumsum(x.cpu(), axis=0)
     assert torch.allclose(res, res2)
+
+
+def _reference_topk(x, dim, k):
+    x = x.cpu().numpy()
+    one_hot = np.zeros((x.shape[0], x.shape[1]), dtype=int)
+    if dim == 1:
+        for i in range(x.shape[0]):
+            one_hot[i, np.argsort(x[i, :])[::-1][:k]] = 1
+    for i in range(x.shape[1]):
+        one_hot[np.argsort(x[:, i])[::-1][:k], i] = 1
+    return one_hot
+
+
+@pytest.mark.parametrize("dtype", [torch.half, torch.float, torch.double])
+@pytest.mark.parametrize("k", [3, 5])
+@pytest.mark.parametrize("dim", [0, 1])
+def test_custom_topk(dtype, k, dim):
+    """Test custom topk implementation."""
+    x = torch.randn(100, 10, dtype=dtype)
+    top_k = select_topk(x, dim=dim, topk=k)
+    assert top_k.shape == (100, 10)
+    assert top_k.dtype == torch.int
+    ref = _reference_topk(x, dim=dim, k=k)
+    assert torch.allclose(top_k, torch.from_numpy(ref).to(torch.int))
+
+
+def test_half_precision_top_k_cpu_raises_error():
+    """Test that half precision topk raises error on cpu.
+
+    If this begins to fail, it means newer Pytorch versions support this and we can drop internal support.
+
+    """
+    x = torch.randn(100, 10, dtype=torch.half)
+    with pytest.raises(RuntimeError, match="\"topk_cpu\" not implemented for 'Half'"):
+        torch.topk(x, k=3, dim=1)