tfp.substrates.jax.stats.cumulative_variance

Cumulative estimates of variance.

tfp.substrates.jax.stats.cumulative_variance(
    x, sample_axis=0, name=None
)

Given N samples of a scalar-valued random variable X, we can compute cumulative variance estimates

result[i] = variance(x[0:i+1])

in O(N) work and O(log(N)) depth (the length of the longest series of operations that are performed sequentially), with O(1) TF kernel invocations. This implementation also arranges to do so in a numerically accurate manner, i.e., without incurring a subtraction of floating-point numbers of size quadratic in the data x. The underlying algorithm is from [1].

x A numeric Tensor holding samples.
sample_axis Scalar Tensor designating the axis holding samples. Other axes are treated in batch. Default value: 0 (leftmost dimension).
name Python str name prefixed to Ops created by this function. Default value: None (i.e., 'cumulative_variance').

cum_var A Tensor of same shape and dtype as x giving cumulative variance estimates. The zeroth element is the variance of a size-1 set of samples, so 0.

References

[1]: Philippe Pebay. Formulas for Robust, One-Pass Parallel Computation of Covariances and Arbitrary-Order Statistical Moments. Technical Report SAND2008-6212, 2008. https://prod-ng.sandia.gov/techlib-noauth/access-control.cgi/2008/086212.pdf