def testRunningMeanMaxPoints(self):
window_size = 100
rng = np.random.RandomState(_test_seed())
data = tf.convert_to_tensor(
np.concatenate([
rng.randn(window_size), 1. + 2. * rng.randn(window_size * 10)
],
axis=0))
def kernel(rms, idx):
rms, _ = fun_mcmc.running_mean_step(rms,
data[idx],
window_size=window_size)
return (rms, idx + 1), rms.mean
_, mean = fun_mcmc.trace(
state=(fun_mcmc.running_mean_init([], data.dtype), 0),
fn=kernel,
num_steps=len(data),
)
# Up to window_size, we compute the running mean exactly.
self.assertAllClose(np.mean(data[:window_size]), mean[window_size - 1])
# After window_size, we're doing exponential moving average, and pick up the
# mean after the change in the distribution. Since the moving average is
# computed only over ~window_size points, this test is rather noisy.
self.assertAllClose(1., mean[-1], atol=0.2)
def SanitizedAutoCorrelationMean(x,
axis,
reduce_axis,
max_lags=None,
**kwargs):
shape_arr = np.array(list(x.shape))
axes = list(sorted(set(range(len(shape_arr))) - set([reduce_axis])))
mean_shape = shape_arr[axes]
if max_lags is not None:
mean_shape[axis] = max_lags + 1
mean_state = fun_mcmc.running_mean_init(mean_shape, x.dtype)
new_order = list(range(len(shape_arr)))
new_order[0] = new_order[reduce_axis]
new_order[reduce_axis] = 0
x = tf.transpose(x, new_order)
x_arr = tf.TensorArray(x.dtype, x.shape[0]).unstack(x)
mean_state, _ = fun_mcmc.trace(
state=mean_state,
fn=lambda state: fun_mcmc.running_mean_step( # pylint: disable=g-long-lambda
state,
SanitizedAutoCorrelation(x_arr.read(state.num_points),
axis,
max_lags=max_lags,
**kwargs)),
num_steps=x.shape[0],
trace_fn=lambda *_: ())
return mean_state.mean
def testRunningMean(self, shape, aggregation):
rng = np.random.RandomState(_test_seed())
data = tf.convert_to_tensor(rng.randn(*shape))
def kernel(rms, idx):
rms, _ = fun_mcmc.running_mean_step(rms, data[idx], axis=aggregation)
return (rms, idx + 1), ()
true_aggregation = (0,) + (() if aggregation is None else tuple(
[a + 1 for a in util.flatten_tree(aggregation)]))
true_mean = np.mean(data, true_aggregation)
(rms, _), _ = fun_mcmc.trace(
state=(fun_mcmc.running_mean_init(true_mean.shape, data.dtype), 0),
fn=kernel,
num_steps=len(data),
trace_fn=lambda *args: ())
self.assertAllClose(true_mean, rms.mean)
def testSimpleDualAverages(self):
def loss_fn(x, y):
return tf.square(x - 1.) + tf.square(y - 2.), []
def kernel(sda_state, rms_state):
sda_state, _ = fun_mcmc.simple_dual_averages_step(sda_state, loss_fn, 1.)
rms_state, _ = fun_mcmc.running_mean_step(rms_state, sda_state.state)
return (sda_state, rms_state), rms_state.mean
_, (x, y) = fun_mcmc.trace(
(
fun_mcmc.simple_dual_averages_init(
[self._constant(0.), self._constant(0.)]),
fun_mcmc.running_mean_init([[], []], [self._dtype, self._dtype]),
),
kernel,
num_steps=1000,
)
self.assertAllClose(1., x[-1], atol=1e-1)
self.assertAllClose(2., y[-1], atol=1e-1)
