Update tensorflow_probability/python/distributions/categorical.py

Updated Categorical Distribution to use xlogy instead of a boolean mask. Co-Authored-By: harrismirza <[email protected]>
SiegeLordEx · Dec 4, 2018 · eaf626a · eaf626a
1 parent ee50a4c
commit eaf626a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorflow_probability/python/distributions/categorical.py b/tensorflow_probability/python/distributions/categorical.py
@@ -300,7 +300,7 @@ def _log_prob(self, k):
 
   def _entropy(self):
     entropies = self.probs * tf.log(self.probs)
-    return -tf.reduce_sum(tf.boolean_mask(entropies, tf.is_finite(entropies)), axis=-1)
+    return -tf.reduce_sum(tf.xlogy(self.probs, self.probs), axis=-1)
 
   def _mode(self):
     ret = tf.argmax(self.logits, axis=self._batch_rank)