Ensure axes passed to reduce_max or unique in logeinsum_exp.

PiperOrigin-RevId: 319296408
ogrisel · Jul 1, 2020 · 14dd186 · 14dd186
1 parent 02b86d2
commit 14dd186
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tensorflow_probability/python/experimental/marginalize/logeinsumexp.py b/tensorflow_probability/python/experimental/marginalize/logeinsumexp.py
@@ -232,7 +232,7 @@ def _binary_einslogsumexp(formula, a, b):
   # That method can result in -inf entries even though all of the
   # numbers are well within a domain in which `logeinsumexp` shouldn't
   # underflow.
-  almost = tf.reduce_logsumexp(a + b,
-                               axis=tuple(map(intermediates.index, reductions)))
+  almost = tf.reduce_logsumexp(
+      a + b, axis=tuple(set(map(intermediates.index, reductions))))
 
   return rearrange(reduced, rhs, almost)