Add tests to check for nans with logits (pyro-ppl#454)

Author: neerajprad

pyro/distributions/categorical.py

@@ -19,8 +19,9 @@ class Categorical(Distribution):
     :param ps: Probabilities. These should be non-negative and normalized
         along the rightmost axis.
     :type ps: `torch.autograd.Variable`.
-    :param logits: Non-normalized log probability values. Either `ps` or `logits`
-        should be specified but not both.
+    :param logits: Log probability values. When exonentiated, these should
+        sum to 1 along the last axis. Either `ps` or `logits` should be
+        specified but not both.
     :type logits: `torch.autograd.Variable`.
     :param vs: Optional list of values in the support.
     :type vs: `list` or `numpy.array` or `torch.autograd.Variable`

pyro/distributions/util.py

@@ -150,18 +150,20 @@ def _get_clamping_buffer(tensor):
 
 def get_probs_and_logits(ps=None, logits=None, is_multidimensional=True):
     """
-    Convert probability values to logits, or vice-versa. Either `ps` or
-    `logits` should be specified, but not both.
+    Convert probability values to logits, or vice-versa. Either ``ps`` or
+    ``logits`` should be specified, but not both.
 
     :param ps: tensor of probabilities. Should be in the interval *[0, 1]*.
-        If, `is_multidimensional = True`, then must be normalized along
+        If, ``is_multidimensional = True``, then must be normalized along
         axis -1.
-    :param logits: tensor of logit values.
+    :param logits: tensor of logit values.  For the multidimensional case,
+        the values, when exponentiated along the last dimension, must sum
+        to 1.
     :param is_multidimensional: determines the computation of ps from logits,
         and vice-versa. For the multi-dimensional case, logit values are
-        assumed to be non-normalized log probabilities, whereas for the uni-
-        dimensional case, it specifically refers to log odds.
-    :return: tuple containing raw probabilities and logits as tensors
+        assumed to be log probabilities, whereas for the uni-dimensional case,
+        it specifically refers to log odds.
+    :return: tuple containing raw probabilities and logits as tensors.
     """
     assert (ps is None) != (logits is None)
     if ps is not None:

tests/distributions/test_gradient_flow.py

@@ -27,11 +27,41 @@ def test_bernoulli_overflow_gradient(init_tensor_type):
     assert_equal(p.grad.data[0], 0)
 
 
+@pytest.mark.parametrize('init_tensor_type', [torch.FloatTensor])
+def test_bernoulli_with_logits_underflow_gradient(init_tensor_type):
+    p = Variable(init_tensor_type([-1e40]), requires_grad=True)
+    bernoulli = Bernoulli(logits=p)
+    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_with_logits_overflow_gradient(init_tensor_type):
+    p = Variable(init_tensor_type([1e40]), requires_grad=True)
+    bernoulli = Bernoulli(logits=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])))
+    categorical = Categorical(p)
+    log_pdf = categorical.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)
+
+
+@pytest.mark.parametrize('init_tensor_type', [torch.DoubleTensor, torch.FloatTensor])
+def test_categorical_gradient_with_logits(init_tensor_type):
+    p = Variable(init_tensor_type([-float('inf'), 0]), requires_grad=True)
+    categorical = Categorical(logits=p)
+    log_pdf = categorical.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)