tfp.math.reduce_kahan_sum

Reduces the input tensor along the given axis using Kahan summation.

tfp.math.reduce_kahan_sum(
    input_tensor, axis=None, keepdims=False, name=None
)

Returns both the total and the correction term, as a namedtuple, representing the sum in higher precision as total - correction.

A practical use-case is computing the difference of two large (magnitude) sums we expect to be nearly equal. If instead we take their difference as (s0.total - s1.total) - (s0.correction - s1.correction), we can retain more precision in computing their difference.

Note that total holds all the high-order bits of the sum, so the correction can be safely neglected if further enhanced precision computations are not required.

Note: (TF + JAX) This function does not work properly on XLA:CPU without the environment variable: XLA_FLAGS=--xla_cpu_enable_fast_math=false, due to LLVM's reassociation optimizations, which simplify error terms to zero.

Args
input_tensor	The tensor to sum.
axis	One of None, a Python int, or a sequence of Python int. The axes to be reduced. None is taken as "reduce all axes".
keepdims	Python bool indicating whether we return a tensor with singleton dimensions in the reduced axes (True), or squeeze the axes out (default, False).
name	Optional name for ops in scope.

Returns
reduced	A Kahan(total, correction) namedtuple.