Repository: pytorch
Issue: KL divergence between two Continuous Bernoulli is negative
Labels: module: distributions,triaged

### 🐛 Describe the bug

Computing the KL-divergence between two Continuous Bernoulli distributions parametrised by `logits` returns negative values, while it should always be non-negative. Here is a minimal example below

```python
import torch

A = torch.distributions.ContinuousBernoulli(logits=-500)
B = torch.distributions.ContinuousBernoulli(logits=20)

torch.distributions.kl_divergence(A, B) # tensor(-16.6750)
```

## The cause of the issue

The KL is computed as follows:

https://github.com/pytorch/pytorch/blob/7176c92687d3cc847cc046bf002269c6949a21c2/torch/distributions/kl.py#L223-L228

Note that the computation uses both `logits` and `probs` that are derived from `logits` here

https://github.com/pytorch/pytorch/blob/7176c92687d3cc847cc046bf002269c6949a21c2/torch/distributions/continuous_bernoulli.py#L133-L135

However, `probs` are always clamped with `clamp_probs()` to avoid numerical issues. This means that going back from `probs`  to `logits` will not give the same result. This discrepancy between `probs` and `logits` is the reason of the bug.

Parametrising the distribution directly with `probs` argument does not have the same issue:

```python
import torch

A = torch.distributions.ContinuousBernoulli(probs=torch.sigmoid(torch.tensor(-500)))
B = torch.distributions.ContinuousBernoulli(probs=torch.sigmoid(torch.tensor(20)))

torch.distributions.kl_divergence(A, B) # tensor(13.9424)
```
The discrepancy between `logits` and `probs` may cause further bugs too, for example, `log_prob()` uses both logits and probs (via `_cont_bern_log_norm()` method).

## The fix

The fix would be to clamp the `logits` correspondingly to `probs` during initialisation. E.g. if `logits` argument is used, compute the clamped `probs` and then transform them back with the logit function and store them as `self.logits`. 

EDIT: Though this fix looses the numerical advantage of working in the log domain. Perhaps a similar implementation to TFP might be used.

EDIT 2: ~~Seems like the TFP implementation also ~~has numerical issues~~ bugs: https://github.com/tensorflow/probability/issues/1511~~

EDIT 3: The TFP implementation is now fixed, so pytorch could follow it.


### Versions

Collecting environment information...
PyTorch version: 1.10.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Pop!_OS 21.10 (x86_64)
GCC version: (Ubuntu 11.2.0-7ubuntu2) 11.2.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.34

Python version: 3.8.12 | packaged by conda-forge | (default, Jan 30 2022, 23:42:07)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.15-76051515-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3050 Laptop GPU
Nvidia driver version: 470.86
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] pytorch-lightning==1.5.8
[pip3] torch==1.10.1
[pip3] torchmetrics==0.7.1
[pip3] torchvision==0.11.2
[conda] blas                      2.113                       mkl    conda-forge
[conda] blas-devel                3.9.0            13_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1              ha36c431_10    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            13_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            13_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            13_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            13_linux64_mkl    conda-forge
[conda] mkl                       2022.0.1           h8d4b97c_803    conda-forge
[conda] mkl-devel                 2022.0.1           ha770c72_804    conda-forge
[conda] mkl-include               2022.0.1           h8d4b97c_803    conda-forge
[conda] mypy_extensions           0.4.3            py38h578d9bd_4    conda-forge
[conda] numpy                     1.21.2           py38he2449b9_0    conda-forge
[conda] pytorch                   1.10.1          py3.8_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-lightning         1.5.8                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchmetrics              0.7.1                    pypi_0    pypi
[conda] torchvision               0.11.2               py38_cu113    pytorch

cc @fritzo @neerajprad @alicanb @nikitaved