def test_nested_state_dependent_reducers(self):
fake_kernel = test_fixtures.TestTransitionKernel()
nested_reducer = [[
test_fixtures.NaiveMeanReducer(),
test_fixtures.NaiveMeanReducer()
], [test_fixtures.NaiveMeanReducer()]]
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.assertAllEqualNested(
kernel_results.streaming_calculations,
[[[1, 1], [1, 1]], [[1, 1]]],
)
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_reducer_warm_restart(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
result = tfp.experimental.mcmc.sample_chain(
num_results=5,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
)
last_sample, red_res, kernel_results = self.evaluate([
result.final_state, result.reduction_results,
result.final_kernel_results
])
self.assertEqual(3, red_res)
self.assertEqual(5, last_sample)
self.assertEqual(5, kernel_results.counter_1)
self.assertEqual(10, kernel_results.counter_2)
# Warm-restart the underlying kernel and the reduction using the provided
# restart package
result_2 = tfp.experimental.mcmc.sample_chain(num_results=5,
**result.resume_kwargs)
last_sample_2, reduction_result_2, kernel_results_2 = self.evaluate([
result_2.final_state, result_2.reduction_results,
result_2.final_kernel_results
])
self.assertEqual(5.5, reduction_result_2)
self.assertEqual(10, last_sample_2)
self.assertEqual(10, kernel_results_2.counter_1)
self.assertEqual(20, kernel_results_2.counter_2)
def test_simple_reduction(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
result = tfp.experimental.mcmc.sample_chain(
num_results=5,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
)
last_sample, reduction_result, kernel_results = self.evaluate([
result.final_state, result.reduction_results,
result.final_kernel_results
])
self.assertEqual(5, last_sample)
self.assertEqual(3, reduction_result)
self.assertEqual(5, kernel_results.counter_1)
self.assertEqual(10, kernel_results.counter_2)
# Warm-restart the underlying kernel but not the reduction
result_2 = tfp.experimental.mcmc.sample_chain(
num_results=5,
current_state=last_sample,
kernel=fake_kernel,
reducer=fake_reducer,
previous_kernel_results=kernel_results,
)
last_sample_2, reduction_result_2, kernel_results_2 = self.evaluate([
result_2.final_state, result_2.reduction_results,
result_2.final_kernel_results
])
self.assertEqual(10, last_sample_2)
self.assertEqual(8, reduction_result_2)
self.assertEqual(10, kernel_results_2.counter_1)
self.assertEqual(20, kernel_results_2.counter_2)
def test_reducer_warm_restart(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
red_res, last_sample, kr, red_states = tfp.experimental.mcmc.sample_fold(
num_steps=5,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
return_final_reducer_states=True,
)
red_res, last_sample, kernel_results, red_states = self.evaluate(
[red_res, last_sample, kr, red_states])
self.assertEqual(3, red_res)
self.assertEqual(5, last_sample)
self.assertEqual(5, kernel_results.counter_1)
self.assertEqual(10, kernel_results.counter_2)
# Warm-restart the underlying kernel and the reduction
reduction_rslt_2, last_sample_2, kr_2 = tfp.experimental.mcmc.sample_fold(
num_steps=5,
current_state=last_sample,
previous_kernel_results=kernel_results,
kernel=fake_kernel,
reducer=fake_reducer,
previous_reducer_state=red_states)
reduction_rslt_2, last_sample_2, kernel_results_2 = self.evaluate(
[reduction_rslt_2, last_sample_2, kr_2])
self.assertEqual(5.5, reduction_rslt_2)
self.assertEqual(10, last_sample_2)
self.assertEqual(10, kernel_results_2.counter_1)
self.assertEqual(20, kernel_results_2.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_simple_operation(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
seed1, seed2 = tfp.random.split_seed(
test_util.test_seed(sampler_type='stateless'))
reduction_rslt, last_sample, kr = tfp.experimental.mcmc.sample_fold(
num_steps=5,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
seed=seed1)
reduction_rslt, last_sample, kernel_results = self.evaluate(
[reduction_rslt, last_sample, kr])
self.assertEqual(3, reduction_rslt)
self.assertEqual(5, last_sample)
self.assertEqual(5, kernel_results.counter_1)
self.assertEqual(10, kernel_results.counter_2)
# Warm-restart the underlying kernel but not the reduction
reduction_rslt_2, last_sample_2, kr_2 = tfp.experimental.mcmc.sample_fold(
num_steps=5,
current_state=last_sample,
kernel=fake_kernel,
reducer=fake_reducer,
previous_kernel_results=kernel_results,
seed=seed2)
reduction_rslt_2, last_sample_2, kernel_results_2 = self.evaluate(
[reduction_rslt_2, last_sample_2, kr_2])
self.assertEqual(8, reduction_rslt_2)
self.assertEqual(10, last_sample_2)
self.assertEqual(10, kernel_results_2.counter_1)
self.assertEqual(20, kernel_results_2.counter_2)
def test_tracing_a_reduction(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
result = tfp.experimental.mcmc.sample_chain(
num_results=5,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
trace_fn=lambda _state, _kr, reductions: reductions)
trace = self.evaluate(result.trace)
self.assertAllEqual(trace, [1.0, 1.5, 2.0, 2.5, 3.0])
def test_nested_reducers(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducers = [[
test_fixtures.NaiveMeanReducer(),
tfp.experimental.mcmc.CovarianceReducer()
], [test_fixtures.NaiveMeanReducer()]]
reduction_rslt, last_sample, kr = tfp.experimental.mcmc.sample_fold(
num_steps=3,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducers)
reduction_rslt, last_sample, kernel_results = self.evaluate(
[reduction_rslt, last_sample, kr])
self.assertEqual(2, len(reduction_rslt))
self.assertEqual(2, len(reduction_rslt[0]))
self.assertEqual(1, len(reduction_rslt[1]))
self.assertEqual(2, reduction_rslt[0][0])
self.assertNear(2 / 3, reduction_rslt[0][1], err=1e-6)
self.assertEqual(3, last_sample)
self.assertEqual(3, kernel_results.counter_1)
self.assertEqual(6, kernel_results.counter_2)
def test_current_state(self, curr_state):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
reduction_rslt, last_sample, kr = tfp.experimental.mcmc.sample_fold(
num_steps=5,
current_state=curr_state,
kernel=fake_kernel,
reducer=fake_reducer)
reduction_rslt, last_sample, kernel_results = self.evaluate(
[reduction_rslt, last_sample, kr])
self.assertEqual(np.mean(np.arange(curr_state + 1, curr_state + 6)),
reduction_rslt)
self.assertEqual(curr_state + 5, last_sample)
self.assertEqual(5, kernel_results.counter_1)
self.assertEqual(10, kernel_results.counter_2)
def test_simple_operation(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
reduction_rslt, last_sample, kr = tfp.experimental.mcmc.sample_fold(
num_steps=5,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
)
reduction_rslt, last_sample, kernel_results = self.evaluate(
[reduction_rslt, last_sample, kr])
self.assertEqual(3, reduction_rslt)
self.assertEqual(5, last_sample)
self.assertEqual(5, kernel_results.counter_1)
self.assertEqual(10, kernel_results.counter_2)
def test_tensor_thinning_and_burnin(self):
fake_kernel = test_fixtures.TestTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
reduction_rslt, last_sample, kr = tfp.experimental.mcmc.sample_fold(
num_steps=tf.convert_to_tensor(5),
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
num_burnin_steps=tf.convert_to_tensor(10),
num_steps_between_results=tf.convert_to_tensor(1))
reduction_rslt, last_sample, kernel_results = self.evaluate(
[reduction_rslt, last_sample, kr])
self.assertEqual(16, reduction_rslt)
self.assertEqual(20, last_sample)
self.assertEqual(20, kernel_results.counter_1)
self.assertEqual(40, kernel_results.counter_2)
def test_seed_reproducibility(self):
seed = samplers.sanitize_seed(test_util.test_seed())
fake_kernel = test_fixtures.RandomTransitionKernel()
fake_reducer = test_fixtures.NaiveMeanReducer()
first_reduction_rslt, _, _ = tfp.experimental.mcmc.sample_fold(
num_steps=3,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
seed=seed)
second_reduction_rslt, _, _ = tfp.experimental.mcmc.sample_fold(
num_steps=3,
current_state=0.,
kernel=fake_kernel,
reducer=fake_reducer,
seed=seed)
first_reduction_rslt, second_reduction_rslt = self.evaluate(
[first_reduction_rslt, second_reduction_rslt])
self.assertEqual(first_reduction_rslt, second_reduction_rslt)
