def test_is_calibrated(self):
calibrated_kernel = test_fixtures.TestTransitionKernel()
uncalibrated_kernel = test_fixtures.TestTransitionKernel(
is_calibrated=False)
calibrated_thinner = tfp.experimental.mcmc.ThinningKernel(
calibrated_kernel, 0)
uncalibrated_thinner = tfp.experimental.mcmc.ThinningKernel(
uncalibrated_kernel, 0)
self.assertTrue(calibrated_thinner.is_calibrated)
self.assertFalse(uncalibrated_thinner.is_calibrated)
def test_is_calibrated(self):
calibrated_kernel = test_fixtures.TestTransitionKernel()
uncalibrated_kernel = test_fixtures.TestTransitionKernel(
is_calibrated=False)
calibrated_discarder = tfp.experimental.mcmc.SampleDiscardingKernel(
calibrated_kernel)
uncalibrated_discarder = tfp.experimental.mcmc.SampleDiscardingKernel(
uncalibrated_kernel)
self.assertTrue(calibrated_discarder.is_calibrated)
self.assertFalse(uncalibrated_discarder.is_calibrated)
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_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_covariance_after_transformation(self):
fake_kernel = test_fixtures.TestTransitionKernel(
target_log_prob_fn=lambda x: -x**2 / 2)
transformed_kernel = tfp.mcmc.TransformedTransitionKernel(
inner_kernel=fake_kernel,
bijector=tfp.bijectors.Exp(),
)
cov_reducer = tfp.experimental.mcmc.CovarianceReducer()
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=transformed_kernel,
reducer=cov_reducer,
)
samples, _, kernel_results = tfp.mcmc.sample_chain(
num_results=10,
current_state=1.,
kernel=reducer_kernel,
trace_fn=None,
return_final_kernel_results=True,
seed=test_util.test_seed())
samples, final_cov = self.evaluate([
samples,
cov_reducer.finalize(kernel_results.streaming_calculations)
])
self.assertAllClose(
np.mean(samples, axis=0),
kernel_results.streaming_calculations.cov_state.mean,
rtol=1e-6)
self.assertAllClose(np.cov(samples.T, ddof=0), final_cov, rtol=1e-6)
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_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_non_static_thinning(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
num_steps_to_skip = tf.Variable(1, dtype=tf.int32)
thinner = tfp.experimental.mcmc.ThinningKernel(
fake_inner_kernel, num_steps_to_skip=num_steps_to_skip)
def _loop_body(i, curr_state, pkr):
new_state, kernel_results = thinner.one_step(
curr_state,
pkr,
)
return (i + 1, new_state, kernel_results)
pkr = thinner.bootstrap_results(0.)
_, final_sample, kernel_results = tf.while_loop(
lambda i, _, __: i < 2,
_loop_body,
(0., 0., pkr),
)
self.evaluate([num_steps_to_skip.initializer])
final_sample, kernel_results = self.evaluate(
[final_sample, kernel_results])
self.assertEqual(4, final_sample)
self.assertEqual(4, kernel_results.counter_1)
self.assertEqual(8, kernel_results.counter_2)
def test_tensor_thinning(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
thinner = tfp.experimental.mcmc.ThinningKernel(
fake_inner_kernel,
num_steps_to_skip=tf.convert_to_tensor(1),
)
def _loop_body(i, curr_state, pkr):
new_state, kernel_results = thinner.one_step(
curr_state,
pkr,
)
return (i + 1, new_state, kernel_results)
pkr = thinner.bootstrap_results(0.)
_, final_sample, kernel_results = tf.while_loop(
lambda i, _, __: i < 2,
_loop_body,
(0., 0., pkr),
)
final_sample, kernel_results = self.evaluate(
[final_sample, kernel_results])
self.assertEqual(4, final_sample)
self.assertEqual(4, kernel_results.counter_1)
self.assertEqual(8, kernel_results.counter_2)
def test_basic_operation(self):
kernel = test_fixtures.TestTransitionKernel()
result = tfp.experimental.mcmc.sample_chain_with_burnin(
num_results=2,
current_state=0.,
kernel=kernel,
seed=test_util.test_seed())
samples = result.trace
kernel_results = result.final_kernel_results
self.assertAllClose([2], tensorshape_util.as_list(samples.shape))
samples, kernel_results = self.evaluate([samples, kernel_results])
self.assertAllClose([1, 2], samples)
self.assertAllClose(2, kernel_results.counter_1)
self.assertAllClose(4, kernel_results.counter_2)
# Warm-restart the underlying kernel. The Trace does not support warm
# restart.
result_2 = tfp.experimental.mcmc.sample_chain_with_burnin(
num_results=2, **result.resume_kwargs)
samples_2, kernel_results_2 = self.evaluate(
[result_2.trace, result_2.final_kernel_results])
self.assertAllClose([3, 4], samples_2)
self.assertAllClose(4, kernel_results_2.counter_1)
self.assertAllClose(8, kernel_results_2.counter_2)
def test_tf_while(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
thinner = tfp.experimental.mcmc.ThinningKernel(
fake_inner_kernel,
num_steps_to_skip=1,
)
def _loop_body(i, seed, curr_state, pkr):
mcmc_seed, seed = tfp.random.split_seed(seed)
new_state, kernel_results = thinner.one_step(
curr_state,
pkr,
seed=mcmc_seed,
)
return (i + 1, seed, new_state, kernel_results)
pkr = thinner.bootstrap_results(0.)
_, _, final_sample, kernel_results = tf.while_loop(
lambda i, *_: i < 2,
_loop_body,
(0., test_util.test_seed(sampler_type='stateless'), 0., pkr),
)
final_sample, kernel_results = self.evaluate(
[final_sample, kernel_results])
self.assertEqual(4, final_sample)
self.assertEqual(4, kernel_results.counter_1)
self.assertEqual(8, kernel_results.counter_2)
def test_tensor_thinning_and_burnin(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
discarder = tfp.experimental.mcmc.SampleDiscardingKernel(
fake_inner_kernel,
num_burnin_steps=tf.convert_to_tensor(10),
num_steps_between_results=tf.convert_to_tensor(1),)
def _loop_body(i, curr_state, pkr):
new_state, kernel_results = discarder.one_step(
curr_state, pkr,
)
return (i + 1, new_state, kernel_results)
pkr = discarder.bootstrap_results(0.)
_, final_sample, kernel_results = tf.while_loop(
lambda i, _, __: i < 2,
_loop_body,
(0., 0., pkr),
)
final_sample, kernel_results = self.evaluate([
final_sample, kernel_results])
self.assertEqual(14, final_sample)
self.assertEqual(2, kernel_results.call_counter)
self.assertEqual(14, kernel_results.inner_results.counter_1)
self.assertEqual(28, kernel_results.inner_results.counter_2)
def test_with_composed_kernel(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
cov_reducer = tfp.experimental.mcmc.CovarianceReducer()
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=tfp.experimental.mcmc.SampleDiscardingKernel(
inner_kernel=fake_inner_kernel,
num_burnin_steps=10,
num_steps_between_results=2,),
reducer=cov_reducer
)
current_state, kernel_results = 0., reducer_kernel.bootstrap_results(0.)
for _ in range(2):
current_state, kernel_results = reducer_kernel.one_step(
current_state, kernel_results)
current_state, kernel_results, cov = self.evaluate([
current_state,
kernel_results,
cov_reducer.finalize(kernel_results.streaming_calculations),
])
self.assertEqual(16, current_state)
self.assertEqual(2, kernel_results.inner_results.call_counter)
self.assertEqual(
16, kernel_results.inner_results.inner_results.counter_1)
self.assertEqual(
32, kernel_results.inner_results.inner_results.counter_2)
self.assertEqual(np.var([13, 16]), cov)
def test_non_static_thinning_and_burnin(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
num_burnin_steps = tf.Variable(10, dtype=tf.int32)
num_steps_between_results = tf.Variable(1, dtype=tf.int32)
discarder = tfp.experimental.mcmc.SampleDiscardingKernel(
fake_inner_kernel,
num_burnin_steps=num_burnin_steps,
num_steps_between_results=num_steps_between_results)
def _loop_body(i, curr_state, pkr):
new_state, kernel_results = discarder.one_step(
curr_state,
pkr,
)
return (i + 1, new_state, kernel_results)
pkr = discarder.bootstrap_results(0.)
_, final_sample, kernel_results = tf.while_loop(
lambda i, *_: i < 2,
_loop_body,
(0., 0., pkr),
)
self.evaluate([
num_burnin_steps.initializer, num_steps_between_results.initializer
])
final_sample, kernel_results = self.evaluate(
[final_sample, kernel_results])
self.assertEqual(14, final_sample)
self.assertEqual(2, kernel_results.call_counter)
self.assertEqual(14, kernel_results.inner_results.counter_1)
self.assertEqual(28, kernel_results.inner_results.counter_2)
def test_tensor_no_burnin(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
discarder = tfp.experimental.mcmc.SampleDiscardingKernel(
fake_inner_kernel,
num_burnin_steps=tf.convert_to_tensor(0),
num_steps_between_results=tf.convert_to_tensor(1),
)
def _loop_body(i, seed, curr_state, pkr):
mcmc_seed, seed = tfp.random.split_seed(seed)
new_state, kernel_results = discarder.one_step(
curr_state,
pkr,
seed=mcmc_seed,
)
return (i + 1, seed, new_state, kernel_results)
pkr = discarder.bootstrap_results(0.)
_, _, final_sample, kernel_results = tf.while_loop(
lambda i, *_: i < 2,
_loop_body,
(0., test_util.test_seed(sampler_type='stateless'), 0., pkr),
)
final_sample, kernel_results = self.evaluate(
[final_sample, kernel_results])
self.assertEqual(4, final_sample)
self.assertEqual(2, kernel_results.call_counter)
self.assertEqual(4, kernel_results.inner_results.counter_1)
self.assertEqual(8, kernel_results.inner_results.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_with_composed_kernel(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
cov_reducer = tfp.experimental.mcmc.CovarianceReducer()
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=tfp.experimental.mcmc.SampleDiscardingKernel(
inner_kernel=fake_inner_kernel,
num_burnin_steps=10,
num_steps_between_results=2,
),
reducer=cov_reducer)
current_state, kernel_results = 0., reducer_kernel.bootstrap_results(
0.)
seed = test_util.test_seed(sampler_type='stateless')
for _ in range(2):
mcmc_seed, seed = tfp.random.split_seed(seed)
current_state, kernel_results = reducer_kernel.one_step(
current_state, kernel_results, seed=mcmc_seed)
cov = cov_reducer.finalize(kernel_results.reduction_results)
self.assertAllEqual(16, current_state)
self.assertAllEqual(2, kernel_results.inner_results.call_counter)
self.assertAllEqual(
16, kernel_results.inner_results.inner_results.counter_1)
self.assertAllEqual(
32, kernel_results.inner_results.inner_results.counter_2)
self.assertAllEqual(np.var([13, 16]), cov)
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_basic_operation(self):
kernel = test_fixtures.TestTransitionKernel()
samples, kernel_results, final_results = tfp.experimental.mcmc.sample_chain(
num_results=2,
current_state=0.,
kernel=kernel,
return_final_kernel_results=True,
seed=test_util.test_seed())
self.assertAllClose([2], tensorshape_util.as_list(samples.shape))
self.assertAllClose([2],
tensorshape_util.as_list(
kernel_results.counter_1.shape))
self.assertAllClose([2],
tensorshape_util.as_list(
kernel_results.counter_2.shape))
samples, kernel_results = self.evaluate([samples, kernel_results])
self.assertAllClose([1, 2], samples)
self.assertAllClose([1, 2], kernel_results.counter_1)
self.assertAllClose([2, 4], kernel_results.counter_2)
# Warm-restart the underlying kernel. The Trace does not support warm
# restart.
samples_2, kr_2 = tfp.experimental.mcmc.sample_chain(
num_results=2,
current_state=samples[-1],
previous_kernel_results=final_results,
kernel=kernel,
)
samples_2, kernel_results_2 = self.evaluate([samples_2, kr_2])
self.assertAllClose([3, 4], samples_2)
self.assertAllClose([3, 4], kernel_results_2.counter_1)
self.assertAllClose([6, 8], kernel_results_2.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_covariance_with_step_kernel(self):
fake_kernel = test_fixtures.TestTransitionKernel()
cov_reducer = tfp.experimental.mcmc.CovarianceReducer()
reducer_kernel = tfp.experimental.mcmc.WithReductions(
inner_kernel=fake_kernel,
reducer=cov_reducer,
)
chain_state, kernel_results = tfp.experimental.mcmc.step_kernel(
num_steps=6,
current_state=0.,
kernel=reducer_kernel,
return_final_kernel_results=True,
)
chain_state, kernel_results, final_cov = self.evaluate([
chain_state, kernel_results,
cov_reducer.finalize(kernel_results.streaming_calculations)
])
self.assertEqual(6, chain_state)
self.assertEqual(3.5,
kernel_results.streaming_calculations.cov_state.mean)
self.assertNear(np.cov(np.arange(1, 7), ddof=0).tolist(),
final_cov,
err=1e-6)
self.assertEqual(6, kernel_results.inner_results.counter_1)
self.assertEqual(12, kernel_results.inner_results.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_cold_start(self):
fake_inner_kernel = test_fixtures.TestTransitionKernel()
discarder = tfp.experimental.mcmc.SampleDiscardingKernel(
fake_inner_kernel,
num_burnin_steps=10,
num_steps_between_results=1,
)
seed1, seed2 = tfp.random.split_seed(
test_util.test_seed(sampler_type='stateless'))
first_state, _ = discarder.one_step(0.,
discarder.bootstrap_results(0.),
seed=seed1)
second_state, kernel_results = discarder.one_step(
first_state, discarder.bootstrap_results(first_state), seed=seed2)
first_state, second_state, kernel_results = self.evaluate(
[first_state, second_state, kernel_results])
# the first `one_step` performs burn-in and thinning to skip
# `num_burnin_steps` + `num_steps_between_results` samples
self.assertEqual(12, first_state)
# the step count is stored in the kernel results. Cold
# restarting resets the count back to 0 and hence, both
# thinning and burn-in are once again performed.
self.assertEqual(24, second_state)
self.assertEqual(1, kernel_results.call_counter)
self.assertEqual(12, kernel_results.inner_results.counter_1)
self.assertEqual(24, kernel_results.inner_results.counter_2)
def test_tracing_no_reduction(self):
fake_kernel = test_fixtures.TestTransitionKernel()
result = tfp.experimental.mcmc.sample_chain(
num_results=5,
current_state=0.,
kernel=fake_kernel,
trace_fn=lambda state, _kr: state + 10)
trace = self.evaluate(result.trace)
self.assertAllEqual(trace, [11.0, 12.0, 13.0, 14.0, 15.0])
def test_initial_state(self):
fake_kernel = test_fixtures.TestTransitionKernel()
final_state = step_kernel(
num_steps=2,
current_state=1,
kernel=fake_kernel,
)
final_state = self.evaluate(final_state)
self.assertEqual(final_state, 3)
def test_warnings_default(self):
with warnings.catch_warnings(record=True) as triggered:
kernel = test_fixtures.TestTransitionKernel()
tfp.experimental.mcmc.sample_chain(num_results=2,
current_state=0.,
kernel=kernel,
seed=test_util.test_seed())
self.assertTrue(
any('Tracing all kernel results by default is deprecated' in str(
warning.message) for warning in triggered))
def test_no_trace_fn(self):
kernel = test_fixtures.TestTransitionKernel()
samples = tfp.experimental.mcmc.sample_chain(
num_results=2,
current_state=0.,
kernel=kernel,
trace_fn=None,
seed=test_util.test_seed())
self.assertAllClose([2], tensorshape_util.as_list(samples.shape))
samples = self.evaluate(samples)
self.assertAllClose([1, 2], samples)
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_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):
with warnings.catch_warnings(record=True) as triggered:
kernel = test_fixtures.TestTransitionKernel(is_calibrated=False)
tfp.experimental.mcmc.sample_chain(
num_results=2,
current_state=0.,
kernel=kernel,
trace_fn=lambda current_state, kernel_results: kernel_results,
seed=test_util.test_seed())
self.assertTrue(
any('supplied `TransitionKernel` is not calibrated.' in str(
warning.message) for warning in triggered))
