clamp Categorical logit from -inf to min_fifo when calculating entropy (#41002)

mattip · facebook-github-bot · commit a0f110190c2a · 2020-07-14T16:21:12.000-07:00
Summary: Fixes gh-40553 by clamping logit values when calculating Categorical.entropy Pull Request resolved: #41002 Reviewed By: mruberry Differential Revision: D22436432 Pulled By: ngimel fbshipit-source-id: 08b7c7b0c15ab4e5a56b3a8ec0d0237ad360202e
diff --git a/aten/src/ATen/core/ivalue.cpp b/aten/src/ATen/core/ivalue.cpp
@@ -400,7 +400,7 @@ std::ostream& IValue::repr(
     case IValue::Tag::Double: {
       double d = v.toDouble();
       int c = std::fpclassify(d);
-      if (c == FP_NORMAL || c == FP_ZERO) {
+      if ((c == FP_NORMAL || c == FP_ZERO ) && std::abs(d) < 1e10) {
         int64_t i = int64_t(d);
         if (double(i) == d) {
           return out << i << ".";
diff --git a/test/test_distributions.py b/test/test_distributions.py
@@ -1178,6 +1178,11 @@ def ref_log_prob(idx, val, log_prob):
         # check entropy computation
         self.assertEqual(Categorical(p).entropy(), torch.tensor([1.0114, 1.0297]), atol=1e-4, rtol=0)
         self.assertEqual(Categorical(s).entropy(), torch.tensor([0.0, 0.0]))
+        # issue gh-40553
+        logits = p.log()
+        logits[1, 1] = logits[0, 2] = float('-inf')
+        e = Categorical(logits=logits).entropy()
+        self.assertEqual(e, torch.tensor([0.6365, 0.5983]), atol=1e-4, rtol=0)
 
     def test_categorical_enumerate_support(self):
         examples = [
diff --git a/torch/csrc/jit/ir/node_hashing.cpp b/torch/csrc/jit/ir/node_hashing.cpp
@@ -200,7 +200,7 @@ size_t HashNode::operator()(const Node* k) const {
     } else if (
         type->isSubtypeOf(NumberType::get()) &&
         k->kindOf(attr::value) == AttributeKind::f) {
-      constant_hash = std::hash<float>{}(k->f(attr::value));
+      constant_hash = std::hash<double>{}(k->f(attr::value));
     } else if (type->isSubtypeOf(BoolType::get())) {
       constant_hash = std::hash<bool>{}(k->i(attr::value));
     }
diff --git a/torch/distributions/categorical.py b/torch/distributions/categorical.py
@@ -51,6 +51,7 @@ def __init__(self, probs=None, logits=None, validate_args=None):
         else:
             if logits.dim() < 1:
                 raise ValueError("`logits` parameter must be at least one-dimensional.")
+            # Normalize
             self.logits = logits - logits.logsumexp(dim=-1, keepdim=True)
         self._param = self.probs if probs is not None else self.logits
         self._num_events = self._param.size()[-1]
@@ -115,7 +116,9 @@ def log_prob(self, value):
         return log_pmf.gather(-1, value).squeeze(-1)
 
     def entropy(self):
-        p_log_p = self.logits * self.probs
+        min_real = torch.finfo(self.logits.dtype).min
+        logits = torch.clamp(self.logits, min=min_real)
+        p_log_p = logits * self.probs
         return -p_log_p.sum(-1)
 
     def enumerate_support(self, expand=True):

Original file line number	Diff line number	Diff line change
`@@ -200,7 +200,7 @@ size_t HashNode::operator()(const Node* k) const {`
`200`	`200`	`} else if (`
`201`	`201`	`type->isSubtypeOf(NumberType::get()) &&`
`202`	`202`	`k->kindOf(attr::value) == AttributeKind::f) {`
`203`		`- constant_hash = std::hash<float>{}(k->f(attr::value));`
	`203`	`+ constant_hash = std::hash<double>{}(k->f(attr::value));`
`204`	`204`	`} else if (type->isSubtypeOf(BoolType::get())) {`
`205`	`205`	`constant_hash = std::hash<bool>{}(k->i(attr::value));`
`206`	`206`	`}`