Repository: Lightning-AI/torchmetrics

Issue:
KL-Divergence gives NaN on 0-probabilities
## 🐛 Bug

KL-Divergence defined `sum_i p_i log(p_i / q_i)` should allow `p_i`s to be zero, simply ignoring those terms. This is standard behavior, similar to how entropy `sum_i p_i log(1/p_i)` can be computed with some `p_i`s being zero.
(Note `q_i`s should not be allowed 0, unless the corresponding p_i also is.)

### To Reproduce

#### Code sample

```
>>> kl = torchmetrics.KLDivergence()
>>> kl.update(torch.tensor([1,0,0],[1,0,0],[0,0,1]]), torch.tensor(torch.randn(3,3).softmax(dim=-1)))
>>> kl.compute()
tensor(nan)
```

### Expected behavior

Some value other than `nan`.

### Environment

- torchmetrics==0.8.2

### Additional context

Allowing 0 values also makes it easier to use KL-Divergence to compute NLL scores. Similar to how Brier scores can be computed using MeanSquaredError on one-hot vectors.


Discussion:

Resolving patch:
diff --git a/torchmetrics/functional/classification/kl_divergence.py b/torchmetrics/functional/classification/kl_divergence.py
--- a/torchmetrics/functional/classification/kl_divergence.py
+++ b/torchmetrics/functional/classification/kl_divergence.py
@@ -19,7 +19,7 @@
 from typing_extensions import Literal
 
 from torchmetrics.utilities.checks import _check_same_shape
-from torchmetrics.utilities.data import METRIC_EPS
+from torchmetrics.utilities.compute import _safe_xlogy
 
 
 def _kld_update(p: Tensor, q: Tensor, log_prob: bool) -> Tuple[Tensor, int]:
@@ -42,8 +42,7 @@ def _kld_update(p: Tensor, q: Tensor, log_prob: bool) -> Tuple[Tensor, int]:
     else:
         p = p / p.sum(axis=-1, keepdim=True)
         q = q / q.sum(axis=-1, keepdim=True)
-        q = torch.clamp(q, METRIC_EPS)
-        measures = torch.sum(p * torch.log(p / q), axis=-1)
+        measures = _safe_xlogy(p, p / q).sum(axis=-1)
 
     return measures, total
 
diff --git a/torchmetrics/functional/regression/tweedie_deviance.py b/torchmetrics/functional/regression/tweedie_deviance.py
--- a/torchmetrics/functional/regression/tweedie_deviance.py
+++ b/torchmetrics/functional/regression/tweedie_deviance.py
@@ -17,13 +17,7 @@
 from torch import Tensor
 
 from torchmetrics.utilities.checks import _check_same_shape
-
-
-def xlogy(x: Tensor, y: Tensor) -> Tensor:
-    """Computes x * log(y). Returns 0 if x=0"""
-    res = x * torch.log(y)
-    res[x == 0] = 0.0
-    return res
+from torchmetrics.utilities.compute import _safe_xlogy
 
 
 def _tweedie_deviance_score_update(preds: Tensor, targets: Tensor, power: float = 0.0) -> Tuple[Tensor, Tensor]:
@@ -58,7 +52,7 @@ def _tweedie_deviance_score_update(preds: Tensor, targets: Tensor, power: float
                 f"For power={power}, 'preds' has to be strictly positive and 'targets' cannot be negative."
             )
 
-        deviance_score = 2 * (xlogy(targets, targets / preds) + preds - targets)
+        deviance_score = 2 * (_safe_xlogy(targets, targets / preds) + preds - targets)
     elif power == 2:
         # Gamma distribution
         if torch.any(preds <= 0) or torch.any(targets <= 0):
diff --git a/torchmetrics/utilities/compute.py b/torchmetrics/utilities/compute.py
--- a/torchmetrics/utilities/compute.py
+++ b/torchmetrics/utilities/compute.py
@@ -23,3 +23,18 @@ def _safe_matmul(x: Tensor, y: Tensor) -> Tensor:
     if x.dtype == torch.float16 or y.dtype == torch.float16:
         return (x.float() @ y.T.float()).half()
     return x @ y.T
+
+
+def _safe_xlogy(x: Tensor, y: Tensor) -> Tensor:
+    """Computes x * log(y). Returns 0 if x=0.
+
+    Example:
+        >>> import torch
+        >>> x = torch.zeros(1)
+        >>> _safe_xlogy(x, 1/x)
+        tensor([0.])
+
+    """
+    res = x * torch.log(y)
+    res[x == 0] = 0.0
+    return res


Test patch:
diff --git a/tests/classification/test_kl_divergence.py b/tests/classification/test_kl_divergence.py
--- a/tests/classification/test_kl_divergence.py
+++ b/tests/classification/test_kl_divergence.py
@@ -112,3 +112,20 @@ def test_error_on_multidim_tensors():
     metric = KLDivergence()
     with pytest.raises(ValueError, match="Expected both p and q distribution to be 2D but got 3 and 3 respectively"):
         metric(torch.randn(10, 20, 5), torch.randn(10, 20, 5))
+
+
+def test_zero_probability():
+    """When p = 0 in kl divergence the score should not output Nan."""
+    metric = KLDivergence()
+    metric.update(
+        torch.tensor([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
+        torch.tensor(torch.randn(3, 3).softmax(dim=-1)),
+    )
+    assert not torch.isnan(metric.compute())
+
+
+def test_inf_case():
+    """When q = 0 in kl divergence the score should be inf."""
+    metric = KLDivergence()
+    metric.update(torch.tensor([[0.3, 0.3, 0.4]]), torch.tensor([[0.5, 0.5, 0]]))
+    assert not torch.isfinite(metric.compute())