def test_nested_reducers(self):
fake_kernel = test_fixtures.TestTransitionKernel()
nested_reducer = [[
test_fixtures.TestReducer(),
test_fixtures.TestReducer()
], [test_fixtures.TestReducer()]]
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=fake_kernel,
reducer=nested_reducer,
)
pkr = reducer_kernel.bootstrap_results(0.)
new_sample, kernel_results = reducer_kernel.one_step(0., pkr)
new_sample, kernel_results = self.evaluate(
[new_sample, kernel_results])
self.assertEqual(2, len(kernel_results.streaming_calculations[0]))
self.assertEqual(1, len(kernel_results.streaming_calculations[1]))
self.assertEqual((2, ),
np.array(kernel_results.streaming_calculations).shape)
self.assertAllEqual([[1, 1], [1]],
kernel_results.streaming_calculations)
self.assertEqual(1, new_sample)
self.assertEqual(1, kernel_results.inner_results.counter_1)
self.assertEqual(2, kernel_results.inner_results.counter_2)
def test_tf_while(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.TestReducer()
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=fake_kernel,
reducer=fake_reducer,
)
initial_state = 0.
initial_kernel_results = reducer_kernel.bootstrap_results(
initial_state)
def _loop_body(i, curr_state, pkr):
new_state, kernel_results = reducer_kernel.one_step(
curr_state, pkr)
return (i + 1, new_state, kernel_results)
_, new_sample, kernel_results = tf.while_loop(
lambda i, _, __: i < 6, _loop_body,
(0., initial_state, initial_kernel_results))
new_sample, kernel_results = self.evaluate(
[new_sample, kernel_results])
self.assertEqual(6, kernel_results.streaming_calculations)
self.assertEqual(6, new_sample)
self.assertEqual(6, kernel_results.inner_results.counter_1)
self.assertEqual(12, kernel_results.inner_results.counter_2)
def test_nested_reducers(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.TestReducer()
mean_reducer = test_fixtures.NaiveMeanReducer()
cov_reducer = tfp.experimental.mcmc.CovarianceReducer()
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=fake_kernel,
reducer=[[mean_reducer, cov_reducer], [fake_reducer]],
)
chain_state, kernel_results = 0., reducer_kernel.bootstrap_results(0.)
for _ in range(6):
chain_state, kernel_results = reducer_kernel.one_step(
chain_state, kernel_results)
final_cov, final_mean = self.evaluate([
cov_reducer.finalize(kernel_results.streaming_calculations[0][1]),
mean_reducer.finalize(kernel_results.streaming_calculations[0][0])
])
self.assertEqual(2, len(kernel_results.streaming_calculations))
self.assertEqual(2, len(kernel_results.streaming_calculations[0]))
self.assertEqual(1, len(kernel_results.streaming_calculations[1]))
self.assertEqual(3.5, final_mean)
self.assertAllEqual(
3.5, kernel_results.streaming_calculations[0][1].cov_state.mean)
self.assertAllEqual(6, kernel_results.streaming_calculations[1][0])
self.assertNear(np.cov(np.arange(1, 7), ddof=0).tolist(),
final_cov,
err=1e-6)
self.assertAllEqual(6, kernel_results.inner_results.counter_1)
self.assertAllEqual(12, kernel_results.inner_results.counter_2)
def test_boostrap_results(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.TestReducer()
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=fake_kernel,
reducer=fake_reducer,
)
pkr = self.evaluate(reducer_kernel.bootstrap_results(9.))
self.assertEqual(0, pkr.streaming_calculations, 0)
self.assertEqual(0, pkr.inner_results.counter_1, 0)
self.assertEqual(0, pkr.inner_results.counter_2, 0)
def test_is_calibrated(self):
fake_calibrated_kernel = test_fixtures.TestTransitionKernel()
fake_uncalibrated_kernel = test_fixtures.TestTransitionKernel(
is_calibrated=False)
fake_reducer = test_fixtures.TestReducer()
calibrated_reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=fake_calibrated_kernel,
reducer=fake_reducer,
)
uncalibrated_reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=fake_uncalibrated_kernel,
reducer=fake_reducer,
)
self.assertTrue(calibrated_reducer_kernel.is_calibrated)
self.assertFalse(uncalibrated_reducer_kernel.is_calibrated)
def test_simple_operation(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.TestReducer()
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=fake_kernel,
reducer=fake_reducer,
)
pkr = reducer_kernel.bootstrap_results(0., )
new_sample, kernel_results = reducer_kernel.one_step(0., pkr)
new_sample, kernel_results = self.evaluate(
[new_sample, kernel_results])
self.assertEqual(1, kernel_results.streaming_calculations)
self.assertEqual(1, new_sample)
self.assertEqual(1, kernel_results.inner_results.counter_1)
self.assertEqual(2, kernel_results.inner_results.counter_2)
