Repository: pyro-ppl/pyro

Issue:
New Bernoulli parameterization lead to subtle bug in AIR model
When I merged #416 into a local branch it broke the AIR example.

Here's a simplified instance of the problem:

```python
import torch
from torch.autograd import Variable
from pyro.distributions import Bernoulli
from torch.nn.functional import sigmoid

p = Variable(torch.Tensor([0]), requires_grad=True)
b = Bernoulli(sigmoid(p) * 0.0)
log_pdf = b.batch_log_pdf(Variable(torch.Tensor([0])))
log_pdf.sum().backward()
print(p.grad)
```

Prior to #416 this returned a `grad` of zero as expected, but it now returns `nan`:

```
pyro$ git rev-parse --short HEAD
71bca18
pyro$ python3 bern.py 
Variable containing:
-0
[torch.FloatTensor of size 1]

pyro$ git rev-parse --short HEAD
a85525a
pyro$ python3 bern.py 
Variable containing:
nan
[torch.FloatTensor of size 1]
```

I suspect that the problem is that converting between probabilities and log odds introduces an intemediate `-inf`, which messes up autograd.

I may be able to adjust the model to work around this, but either way, should this be considered a bug? (It seems like it could trip other people up, and chasing down the source of the `nan`s is tricky.)



Discussion:
Thanks for the reproducible example, @null-a! I'll try to get a fix in today, failing which, it may be better to revert the logit PR instead. Will keep you posted. 
> Thanks for the reproducible example,

@neerajprad No problem, thanks for taking a look!

Resolving patch:
diff --git a/pyro/distributions/util.py b/pyro/distributions/util.py
--- a/pyro/distributions/util.py
+++ b/pyro/distributions/util.py
@@ -133,6 +133,15 @@ def softmax(x, dim=-1):
     return soft_max_nd.transpose(dim, len(input_size) - 1)
 
 
+def _get_clamping_buffer(tensor):
+    clamp_eps = 1e-6
+    if isinstance(tensor, Variable):
+        tensor = tensor.data
+    if isinstance(tensor, (torch.DoubleTensor, torch.cuda.DoubleTensor)):
+        clamp_eps = 1e-15
+    return clamp_eps
+
+
 def get_probs_and_logits(ps=None, logits=None, is_multidimensional=True):
     """
     Convert probability values to logits, or vice-versa. Either `ps` or
@@ -149,14 +158,17 @@ def get_probs_and_logits(ps=None, logits=None, is_multidimensional=True):
     :return: tuple containing raw probabilities and logits as tensors
     """
     assert (ps is None) != (logits is None)
+    if ps is not None:
+        eps = _get_clamping_buffer(ps)
+        ps_clamped = ps.clamp(min=eps, max=1 - eps)
     if is_multidimensional:
         if ps is None:
             ps = softmax(logits, -1)
         else:
-            logits = torch.log(ps)
+            logits = torch.log(ps_clamped)
     else:
         if ps is None:
             ps = F.sigmoid(logits)
         else:
-            logits = torch.log(ps) - torch.log1p(-ps)
+            logits = torch.log(ps_clamped) - torch.log1p(-ps_clamped)
     return ps, logits


Test patch:
diff --git a/tests/distributions/test_gradient_flow.py b/tests/distributions/test_gradient_flow.py
new file mode 100644
--- /dev/null
+++ b/tests/distributions/test_gradient_flow.py
@@ -0,0 +1,37 @@
+import pytest
+import torch
+from torch.autograd import Variable
+from torch.nn.functional import sigmoid
+
+from pyro.distributions import Bernoulli, Categorical
+from tests.common import assert_equal
+
+
+@pytest.mark.parametrize('init_tensor_type', [torch.DoubleTensor, torch.FloatTensor])
+def test_bernoulli_underflow_gradient(init_tensor_type):
+    p = Variable(init_tensor_type([0]), requires_grad=True)
+    bernoulli = Bernoulli(sigmoid(p) * 0.0)
+    log_pdf = bernoulli.batch_log_pdf(Variable(init_tensor_type([0])))
+    log_pdf.sum().backward()
+    assert_equal(log_pdf.data[0], 0)
+    assert_equal(p.grad.data[0], 0)
+
+
+@pytest.mark.parametrize('init_tensor_type', [torch.DoubleTensor, torch.FloatTensor])
+def test_bernoulli_overflow_gradient(init_tensor_type):
+    p = Variable(init_tensor_type([1e32]), requires_grad=True)
+    bernoulli = Bernoulli(sigmoid(p))
+    log_pdf = bernoulli.batch_log_pdf(Variable(init_tensor_type([1])))
+    log_pdf.sum().backward()
+    assert_equal(log_pdf.data[0], 0)
+    assert_equal(p.grad.data[0], 0)
+
+
+@pytest.mark.parametrize('init_tensor_type', [torch.DoubleTensor, torch.FloatTensor])
+def test_categorical_gradient(init_tensor_type):
+    p = Variable(init_tensor_type([0, 1]), requires_grad=True)
+    bernoulli = Categorical(p)
+    log_pdf = bernoulli.batch_log_pdf(Variable(init_tensor_type([0, 1])))
+    log_pdf.sum().backward()
+    assert_equal(log_pdf.data[0], 0)
+    assert_equal(p.grad.data[0], 0)