Repository: Lightning-AI/torchmetrics

Issue:
Weird behavior of `CosineSimilarityMetric` when used with tensors of shape [d]
## 🐛 Bug?

According to the documentation, `CosineSimilarityMetric` requires tensors of shape `[N, d]`, with `N` the batch size and `d` the dimension of the vectors.

Using it with vectors of shape `[d]` does not raise any error, and the call to `compute` behaves weirdly and gives confusing results. I'm not sure whether this is the expected behavior, a bug, or if this usage is simply unintended (and then maybe an error could have been raised).

### To Reproduce

```python
>>> from torchmetrics import CosineSimilarity
>>> from torch import tensor
>>> 
>>> cosine_similarity = CosineSimilarity(reduction="mean")
>>> a = tensor([1., 1., 1.])
>>> b = tensor([100., 100., 100.])
>>> cosine_similarity(a, b)
tensor(1.)  # a and b have the same direction, so this is normal.
>>> cosine_similarity(b, a)  
tensor(1.)  # same for b and a.
>>> cosine_similarity.compute()
tensor(0.0200)  # I would expect this to be 1 too (the average of the 2 previous calls).
```
The obtained result (0.02) is actually the cosine similarity between `[1, 1, 1, 100, 100, 100]` and `[100, 100, 100, 1, 1, 1]`. I would have expected to get instead the average between the cosine similarity of `[1, 1, 1]` and `[100, 100, 100]` and the cosine similarity of `[100, 100, 100]` and `[1, 1, 1]`, which is 1.

If instead we use it as the documentation says, with tensors of shape `[N, d]`, we get different results:
```python
>>> from torchmetrics import CosineSimilarity
>>> from torch import tensor
>>> 
>>> cosine_similarity = CosineSimilarity(reduction="mean")
>>> a = tensor([[1., 1., 1.]])  # tensor of shape [1, 3] instead of [3]
>>> b = tensor([[100., 100., 100.]])  # tensor of shape [1, 3] instead of [3]
>>> cosine_similarity(a, b)
tensor(1.)
>>> cosine_similarity(b, a)
tensor(1.)
>>> cosine_similarity.compute()
tensor(1.)  # 1 instead of 0.02
```

Environment:
- TorchMetrics 1.2.0
- Python 3.10.10
- torch 2.1.1
- Ubuntu 20.04.6 LTS

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/src/torchmetrics/functional/regression/cosine_similarity.py b/src/torchmetrics/functional/regression/cosine_similarity.py
--- a/src/torchmetrics/functional/regression/cosine_similarity.py
+++ b/src/torchmetrics/functional/regression/cosine_similarity.py
@@ -31,6 +31,11 @@ def _cosine_similarity_update(
 
     """
     _check_same_shape(preds, target)
+    if preds.ndim != 2:
+        raise ValueError(
+            "Expected input to cosine similarity to be 2D tensors of shape `[N,D]` where `N` is the number of samples"
+            f" and `D` is the number of dimensions, but got tensor of shape {preds.shape}"
+        )
     preds = preds.float()
     target = target.float()
 
diff --git a/src/torchmetrics/regression/cosine_similarity.py b/src/torchmetrics/regression/cosine_similarity.py
--- a/src/torchmetrics/regression/cosine_similarity.py
+++ b/src/torchmetrics/regression/cosine_similarity.py
@@ -118,7 +118,7 @@ def plot(
             >>> # Example plotting a single value
             >>> from torchmetrics.regression import CosineSimilarity
             >>> metric = CosineSimilarity()
-            >>> metric.update(randn(10,), randn(10,))
+            >>> metric.update(randn(10,2), randn(10,2))
             >>> fig_, ax_ = metric.plot()
 
         .. plot::
@@ -130,7 +130,7 @@ def plot(
             >>> metric = CosineSimilarity()
             >>> values = []
             >>> for _ in range(10):
-            ...     values.append(metric(randn(10,), randn(10,)))
+            ...     values.append(metric(randn(10,2), randn(10,2)))
             >>> fig, ax = metric.plot(values)
 
         """


Test patch:
diff --git a/tests/unittests/regression/test_cosine_similarity.py b/tests/unittests/regression/test_cosine_similarity.py
--- a/tests/unittests/regression/test_cosine_similarity.py
+++ b/tests/unittests/regression/test_cosine_similarity.py
@@ -30,8 +30,8 @@
 
 
 _single_target_inputs = _Input(
-    preds=torch.rand(NUM_BATCHES, BATCH_SIZE),
-    target=torch.rand(NUM_BATCHES, BATCH_SIZE),
+    preds=torch.rand(NUM_BATCHES, BATCH_SIZE, 1),
+    target=torch.rand(NUM_BATCHES, BATCH_SIZE, 1),
 )
 
 _multi_target_inputs = _Input(
@@ -40,23 +40,10 @@
 )
 
 
-def _multi_target_ref_metric(preds, target, reduction, sk_fn=sk_cosine):
-    sk_preds = preds.view(-1, num_targets).numpy()
-    sk_target = target.view(-1, num_targets).numpy()
-    result_array = sk_fn(sk_target, sk_preds)
-    col = np.diagonal(result_array)
-    col_sum = col.sum()
-    if reduction == "sum":
-        return col_sum
-    if reduction == "mean":
-        return col_sum / len(col)
-    return col
-
-
-def _single_target_ref_metric(preds, target, reduction, sk_fn=sk_cosine):
-    sk_preds = preds.view(-1).numpy()
-    sk_target = target.view(-1).numpy()
-    result_array = sk_fn(np.expand_dims(sk_preds, axis=0), np.expand_dims(sk_target, axis=0))
+def _ref_metric(preds, target, reduction):
+    sk_preds = preds.numpy()
+    sk_target = target.numpy()
+    result_array = sk_cosine(sk_target, sk_preds)
     col = np.diagonal(result_array)
     col_sum = col.sum()
     if reduction == "sum":
@@ -70,8 +57,8 @@ def _single_target_ref_metric(preds, target, reduction, sk_fn=sk_cosine):
 @pytest.mark.parametrize(
     "preds, target, ref_metric",
     [
-        (_single_target_inputs.preds, _single_target_inputs.target, _single_target_ref_metric),
-        (_multi_target_inputs.preds, _multi_target_inputs.target, _multi_target_ref_metric),
+        (_single_target_inputs.preds, _single_target_inputs.target, _ref_metric),
+        (_multi_target_inputs.preds, _multi_target_inputs.target, _ref_metric),
     ],
 )
 class TestCosineSimilarity(MetricTester):
@@ -104,4 +91,11 @@ def test_error_on_different_shape(metric_class=CosineSimilarity):
     """Test that error is raised on different shapes of input."""
     metric = metric_class()
     with pytest.raises(RuntimeError, match="Predictions and targets are expected to have the same shape"):
-        metric(torch.randn(100), torch.randn(50))
+        metric(torch.randn(100, 2), torch.randn(50, 2))
+
+
+def test_error_on_non_2d_input():
+    """Test that error is raised if input is not 2-dimensional."""
+    metric = CosineSimilarity()
+    with pytest.raises(ValueError, match="Expected input to cosine similarity to be 2D tensors of shape.*"):
+        metric(torch.randn(100), torch.randn(100))
diff --git a/tests/unittests/utilities/test_plot.py b/tests/unittests/utilities/test_plot.py
--- a/tests/unittests/utilities/test_plot.py
+++ b/tests/unittests/utilities/test_plot.py
@@ -480,7 +480,7 @@
             id="learned perceptual image patch similarity",
         ),
         pytest.param(ConcordanceCorrCoef, _rand_input, _rand_input, id="concordance corr coef"),
-        pytest.param(CosineSimilarity, _rand_input, _rand_input, id="cosine similarity"),
+        pytest.param(CosineSimilarity, _multilabel_rand_input, _multilabel_rand_input, id="cosine similarity"),
         pytest.param(ExplainedVariance, _rand_input, _rand_input, id="explained variance"),
         pytest.param(KendallRankCorrCoef, _rand_input, _rand_input, id="kendall rank corr coef"),
         pytest.param(