tfp.experimental.marginalize.logeinsumexp

Computing einsum in logarithmic space.

formula a formula compatible with einsum. (But note that ellipses are not supported.)
*tensors tensors to which contractions will be applied.

logeinsumexp(formula, *tensors) is equivalent to tf.math.log(einsum(formula, *map(tf.math.exp, tensors))) except that it is more numerically stable.

This implementation of logeinsumexp is primarily intended for internal use in the marginalize library and assumes the formula is well-formed.