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)