def _kernel_leaves_target_invariant(self,
initial_draws,
event_dims,
sess,
feed_dict=None):
def log_gamma_log_prob(x):
return self._log_gamma_log_prob(x, event_dims)
def fake_log_prob(x):
"""Cooled version of the target distribution."""
return 1.1 * log_gamma_log_prob(x)
step_size = array_ops.placeholder(np.float32, [], name='step_size')
if feed_dict is None:
feed_dict = {}
feed_dict[step_size] = 0.4
sample, acceptance_probs, _, _ = hmc.kernel(step_size, 5,
initial_draws,
log_gamma_log_prob,
event_dims)
bad_sample, bad_acceptance_probs, _, _ = hmc.kernel(
step_size, 5, initial_draws, fake_log_prob, event_dims)
(acceptance_probs_val, bad_acceptance_probs_val, initial_draws_val,
updated_draws_val, fake_draws_val) = sess.run([
acceptance_probs, bad_acceptance_probs, initial_draws, sample,
bad_sample
], feed_dict)
# Confirm step size is small enough that we usually accept.
self.assertGreater(acceptance_probs_val.mean(), 0.5)
self.assertGreater(bad_acceptance_probs_val.mean(), 0.5)
# Confirm step size is large enough that we sometimes reject.
self.assertLess(acceptance_probs_val.mean(), 0.99)
self.assertLess(bad_acceptance_probs_val.mean(), 0.99)
_, ks_p_value_true = stats.ks_2samp(initial_draws_val.flatten(),
updated_draws_val.flatten())
_, ks_p_value_fake = stats.ks_2samp(initial_draws_val.flatten(),
fake_draws_val.flatten())
logging.vlog(
1, 'acceptance rate for true target: {}'.format(
acceptance_probs_val.mean()))
logging.vlog(
1, 'acceptance rate for fake target: {}'.format(
bad_acceptance_probs_val.mean()))
logging.vlog(1,
'K-S p-value for true target: {}'.format(ks_p_value_true))
logging.vlog(1,
'K-S p-value for fake target: {}'.format(ks_p_value_fake))
# Make sure that the MCMC update hasn't changed the empirical CDF much.
self.assertGreater(ks_p_value_true, 1e-3)
# Confirm that targeting the wrong distribution does
# significantly change the empirical CDF.
self.assertLess(ks_p_value_fake, 1e-6)
def testNanFromGradsDontPropagate(self):
"""Test that update with NaN gradients does not cause NaN in results."""
def _nan_log_prob_with_nan_gradient(x):
return np.nan * math_ops.reduce_sum(x)
with self.test_session() as sess:
initial_x = math_ops.linspace(0.01, 5, 10)
updated_x, kernel_results = hmc.kernel(
target_log_prob_fn=_nan_log_prob_with_nan_gradient,
current_state=initial_x,
step_size=2.,
num_leapfrog_steps=5,
seed=47)
initial_x_, updated_x_, acceptance_probs_ = sess.run(
[initial_x, updated_x, kernel_results.acceptance_probs])
logging_ops.vlog(1, "initial_x = {}".format(initial_x_))
logging_ops.vlog(1, "updated_x = {}".format(updated_x_))
logging_ops.vlog(1, "acceptance_probs = {}".format(acceptance_probs_))
self.assertAllEqual(initial_x_, updated_x_)
self.assertEqual(acceptance_probs_, 0.)
self.assertAllFinite(
gradients_ops.gradients(updated_x, initial_x)[0].eval())
self.assertAllEqual([True], [g is None for g in gradients_ops.gradients(
kernel_results.proposed_grads_target_log_prob, initial_x)])
self.assertAllEqual([False], [g is None for g in gradients_ops.gradients(
kernel_results.proposed_grads_target_log_prob,
kernel_results.proposed_state)])
def testNanRejection(self):
"""Tests that an update that yields NaN potentials gets rejected.
We run HMC with a target distribution that returns NaN
log-likelihoods if any element of x < 0, and unit-scale
exponential log-likelihoods otherwise. The exponential potential
pushes x towards 0, ensuring that any reasonably large update will
push us over the edge into NaN territory.
"""
def _unbounded_exponential_log_prob(x):
"""An exponential distribution with log-likelihood NaN for x < 0."""
per_element_potentials = array_ops.where(
x < 0.,
array_ops.fill(array_ops.shape(x), x.dtype.as_numpy_dtype(np.nan)),
-x)
return math_ops.reduce_sum(per_element_potentials)
with self.test_session() as sess:
initial_x = math_ops.linspace(0.01, 5, 10)
updated_x, kernel_results = hmc.kernel(
target_log_prob_fn=_unbounded_exponential_log_prob,
current_state=initial_x,
step_size=2.,
num_leapfrog_steps=5,
seed=46)
initial_x_, updated_x_, acceptance_probs_ = sess.run(
[initial_x, updated_x, kernel_results.acceptance_probs])
logging_ops.vlog(1, "initial_x = {}".format(initial_x_))
logging_ops.vlog(1, "updated_x = {}".format(updated_x_))
logging_ops.vlog(1, "acceptance_probs = {}".format(acceptance_probs_))
self.assertAllEqual(initial_x_, updated_x_)
self.assertEqual(acceptance_probs_, 0.)
def testNanFromGradsDontPropagate(self):
"""Test that update with NaN gradients does not cause NaN in results."""
def _nan_log_prob_with_nan_gradient(x):
return np.nan * math_ops.reduce_sum(x)
with self.test_session(graph=ops.Graph()) as sess:
initial_x = math_ops.linspace(0.01, 5, 10)
updated_x, kernel_results = hmc.kernel(
target_log_prob_fn=_nan_log_prob_with_nan_gradient,
current_state=initial_x,
step_size=2.,
num_leapfrog_steps=5,
seed=47)
initial_x_, updated_x_, acceptance_probs_ = sess.run(
[initial_x, updated_x, kernel_results.acceptance_probs])
logging_ops.vlog(1, "initial_x = {}".format(initial_x_))
logging_ops.vlog(1, "updated_x = {}".format(updated_x_))
logging_ops.vlog(1,
"acceptance_probs = {}".format(acceptance_probs_))
self.assertAllEqual(initial_x_, updated_x_)
self.assertEqual(acceptance_probs_, 0.)
self.assertAllFinite(
gradients_ops.gradients(updated_x, initial_x)[0].eval())
self.assertAllEqual([True], [
g is None for g in gradients_ops.gradients(
kernel_results.proposed_grads_target_log_prob, initial_x)
])
self.assertAllEqual([False], [
g is None for g in gradients_ops.gradients(
kernel_results.proposed_grads_target_log_prob,
kernel_results.proposed_state)
])
def testNanRejection(self):
"""Tests that an update that yields NaN potentials gets rejected.
We run HMC with a target distribution that returns NaN
log-likelihoods if any element of x < 0, and unit-scale
exponential log-likelihoods otherwise. The exponential potential
pushes x towards 0, ensuring that any reasonably large update will
push us over the edge into NaN territory.
"""
def _unbounded_exponential_log_prob(x):
"""An exponential distribution with log-likelihood NaN for x < 0."""
per_element_potentials = array_ops.where(
x < 0.,
array_ops.fill(array_ops.shape(x),
x.dtype.as_numpy_dtype(np.nan)), -x)
return math_ops.reduce_sum(per_element_potentials)
with self.test_session(graph=ops.Graph()) as sess:
initial_x = math_ops.linspace(0.01, 5, 10)
updated_x, kernel_results = hmc.kernel(
target_log_prob_fn=_unbounded_exponential_log_prob,
current_state=initial_x,
step_size=2.,
num_leapfrog_steps=5,
seed=46)
initial_x_, updated_x_, acceptance_probs_ = sess.run(
[initial_x, updated_x, kernel_results.acceptance_probs])
logging_ops.vlog(1, "initial_x = {}".format(initial_x_))
logging_ops.vlog(1, "updated_x = {}".format(updated_x_))
logging_ops.vlog(1,
"acceptance_probs = {}".format(acceptance_probs_))
self.assertAllEqual(initial_x_, updated_x_)
self.assertEqual(acceptance_probs_, 0.)
def testNanFromGradsDontPropagate(self):
"""Test that update with NaN gradients does not cause NaN in results."""
def _nan_log_prob_with_nan_gradient(x):
return np.nan * math_ops.reduce_sum(x)
with self.test_session() as sess:
initial_x = math_ops.linspace(0.01, 5, 10)
updated_x, acceptance_probs, new_log_prob, new_grad = hmc.kernel(
2., 5, initial_x, _nan_log_prob_with_nan_gradient, [0])
initial_x_val, updated_x_val, acceptance_probs_val = sess.run(
[initial_x, updated_x, acceptance_probs])
logging.vlog(1, 'initial_x = {}'.format(initial_x_val))
logging.vlog(1, 'updated_x = {}'.format(updated_x_val))
logging.vlog(1,
'acceptance_probs = {}'.format(acceptance_probs_val))
self.assertAllEqual(initial_x_val, updated_x_val)
self.assertEqual(acceptance_probs_val, 0.)
self.assertAllFinite(
gradients_impl.gradients(updated_x, initial_x)[0].eval())
self.assertTrue(
gradients_impl.gradients(new_grad, initial_x)[0] is None)
# Gradients of the acceptance probs and new log prob are not finite.
_ = new_log_prob # Prevent unused arg error.
def testNanRejection(self):
"""Tests that an update that yields NaN potentials gets rejected.
We run HMC with a target distribution that returns NaN
log-likelihoods if any element of x < 0, and unit-scale
exponential log-likelihoods otherwise. The exponential potential
pushes x towards 0, ensuring that any reasonably large update will
push us over the edge into NaN territory.
"""
def _unbounded_exponential_log_prob(x):
"""An exponential distribution with log-likelihood NaN for x < 0."""
per_element_potentials = array_ops.where(
x < 0, np.nan * array_ops.ones_like(x), -x)
return math_ops.reduce_sum(per_element_potentials)
with self.test_session() as sess:
initial_x = math_ops.linspace(0.01, 5, 10)
updated_x, acceptance_probs, _, _ = hmc.kernel(
2., 5, initial_x, _unbounded_exponential_log_prob, [0])
initial_x_val, updated_x_val, acceptance_probs_val = sess.run(
[initial_x, updated_x, acceptance_probs])
logging.vlog(1, 'initial_x = {}'.format(initial_x_val))
logging.vlog(1, 'updated_x = {}'.format(updated_x_val))
logging.vlog(1,
'acceptance_probs = {}'.format(acceptance_probs_val))
self.assertAllEqual(initial_x_val, updated_x_val)
self.assertEqual(acceptance_probs_val, 0.)
def testNanFromGradsDontPropagate(self):
"""Test that update with NaN gradients does not cause NaN in results."""
def _nan_log_prob_with_nan_gradient(x):
return np.nan * math_ops.reduce_sum(x)
with self.test_session() as sess:
initial_x = math_ops.linspace(0.01, 5, 10)
updated_x, acceptance_probs, new_log_prob, new_grad = hmc.kernel(
2., 5, initial_x, _nan_log_prob_with_nan_gradient, [0])
initial_x_val, updated_x_val, acceptance_probs_val = sess.run(
[initial_x, updated_x, acceptance_probs])
logging.vlog(1, 'initial_x = {}'.format(initial_x_val))
logging.vlog(1, 'updated_x = {}'.format(updated_x_val))
logging.vlog(1, 'acceptance_probs = {}'.format(acceptance_probs_val))
self.assertAllEqual(initial_x_val, updated_x_val)
self.assertEqual(acceptance_probs_val, 0.)
self.assertAllFinite(
gradients_impl.gradients(updated_x, initial_x)[0].eval())
self.assertTrue(
gradients_impl.gradients(new_grad, initial_x)[0] is None)
# Gradients of the acceptance probs and new log prob are not finite.
_ = new_log_prob # Prevent unused arg error.
def testNanRejection(self):
"""Tests that an update that yields NaN potentials gets rejected.
We run HMC with a target distribution that returns NaN
log-likelihoods if any element of x < 0, and unit-scale
exponential log-likelihoods otherwise. The exponential potential
pushes x towards 0, ensuring that any reasonably large update will
push us over the edge into NaN territory.
"""
def _unbounded_exponential_log_prob(x):
"""An exponential distribution with log-likelihood NaN for x < 0."""
per_element_potentials = array_ops.where(x < 0,
np.nan * array_ops.ones_like(x),
-x)
return math_ops.reduce_sum(per_element_potentials)
with self.test_session() as sess:
initial_x = math_ops.linspace(0.01, 5, 10)
updated_x, acceptance_probs, _, _ = hmc.kernel(
2., 5, initial_x, _unbounded_exponential_log_prob, [0])
initial_x_val, updated_x_val, acceptance_probs_val = sess.run(
[initial_x, updated_x, acceptance_probs])
logging.vlog(1, 'initial_x = {}'.format(initial_x_val))
logging.vlog(1, 'updated_x = {}'.format(updated_x_val))
logging.vlog(1, 'acceptance_probs = {}'.format(acceptance_probs_val))
self.assertAllEqual(initial_x_val, updated_x_val)
self.assertEqual(acceptance_probs_val, 0.)
def _kernel_leaves_target_invariant(self, initial_draws, event_dims,
sess, feed_dict=None):
def log_gamma_log_prob(x):
return self._log_gamma_log_prob(x, event_dims)
def fake_log_prob(x):
"""Cooled version of the target distribution."""
return 1.1 * log_gamma_log_prob(x)
step_size = array_ops.placeholder(np.float32, [], name='step_size')
if feed_dict is None:
feed_dict = {}
feed_dict[step_size] = 0.4
sample, acceptance_probs, _, _ = hmc.kernel(step_size, 5, initial_draws,
log_gamma_log_prob, event_dims)
bad_sample, bad_acceptance_probs, _, _ = hmc.kernel(
step_size, 5, initial_draws, fake_log_prob, event_dims)
(acceptance_probs_val, bad_acceptance_probs_val, initial_draws_val,
updated_draws_val, fake_draws_val) = sess.run([acceptance_probs,
bad_acceptance_probs,
initial_draws, sample,
bad_sample], feed_dict)
# Confirm step size is small enough that we usually accept.
self.assertGreater(acceptance_probs_val.mean(), 0.5)
self.assertGreater(bad_acceptance_probs_val.mean(), 0.5)
# Confirm step size is large enough that we sometimes reject.
self.assertLess(acceptance_probs_val.mean(), 0.99)
self.assertLess(bad_acceptance_probs_val.mean(), 0.99)
_, ks_p_value_true = stats.ks_2samp(initial_draws_val.flatten(),
updated_draws_val.flatten())
_, ks_p_value_fake = stats.ks_2samp(initial_draws_val.flatten(),
fake_draws_val.flatten())
logging.vlog(1, 'acceptance rate for true target: {}'.format(
acceptance_probs_val.mean()))
logging.vlog(1, 'acceptance rate for fake target: {}'.format(
bad_acceptance_probs_val.mean()))
logging.vlog(1, 'K-S p-value for true target: {}'.format(ks_p_value_true))
logging.vlog(1, 'K-S p-value for fake target: {}'.format(ks_p_value_fake))
# Make sure that the MCMC update hasn't changed the empirical CDF much.
self.assertGreater(ks_p_value_true, 1e-3)
# Confirm that targeting the wrong distribution does
# significantly change the empirical CDF.
self.assertLess(ks_p_value_fake, 1e-6)
def _kernel_leaves_target_invariant(self, initial_draws,
independent_chain_ndims,
sess, feed_dict=None):
def log_gamma_log_prob(x):
event_dims = math_ops.range(independent_chain_ndims, array_ops.rank(x))
return self._log_gamma_log_prob(x, event_dims)
def fake_log_prob(x):
"""Cooled version of the target distribution."""
return 1.1 * log_gamma_log_prob(x)
step_size = array_ops.placeholder(np.float32, [], name="step_size")
if feed_dict is None:
feed_dict = {}
feed_dict[step_size] = 0.4
sample, kernel_results = hmc.kernel(
target_log_prob_fn=log_gamma_log_prob,
current_state=initial_draws,
step_size=step_size,
num_leapfrog_steps=5,
seed=43)
bad_sample, bad_kernel_results = hmc.kernel(
target_log_prob_fn=fake_log_prob,
current_state=initial_draws,
step_size=step_size,
num_leapfrog_steps=5,
seed=44)
[
acceptance_probs_,
bad_acceptance_probs_,
initial_draws_,
updated_draws_,
fake_draws_,
] = sess.run([
kernel_results.acceptance_probs,
bad_kernel_results.acceptance_probs,
initial_draws,
sample,
bad_sample,
], feed_dict)
# Confirm step size is small enough that we usually accept.
self.assertGreater(acceptance_probs_.mean(), 0.5)
self.assertGreater(bad_acceptance_probs_.mean(), 0.5)
# Confirm step size is large enough that we sometimes reject.
self.assertLess(acceptance_probs_.mean(), 0.99)
self.assertLess(bad_acceptance_probs_.mean(), 0.99)
_, ks_p_value_true = stats.ks_2samp(initial_draws_.flatten(),
updated_draws_.flatten())
_, ks_p_value_fake = stats.ks_2samp(initial_draws_.flatten(),
fake_draws_.flatten())
logging_ops.vlog(1, "acceptance rate for true target: {}".format(
acceptance_probs_.mean()))
logging_ops.vlog(1, "acceptance rate for fake target: {}".format(
bad_acceptance_probs_.mean()))
logging_ops.vlog(1, "K-S p-value for true target: {}".format(
ks_p_value_true))
logging_ops.vlog(1, "K-S p-value for fake target: {}".format(
ks_p_value_fake))
# Make sure that the MCMC update hasn't changed the empirical CDF much.
self.assertGreater(ks_p_value_true, 1e-3)
# Confirm that targeting the wrong distribution does
# significantly change the empirical CDF.
self.assertLess(ks_p_value_fake, 1e-6)
def _kernel_leaves_target_invariant(self,
initial_draws,
independent_chain_ndims,
sess,
feed_dict=None):
def log_gamma_log_prob(x):
event_dims = math_ops.range(independent_chain_ndims,
array_ops.rank(x))
return self._log_gamma_log_prob(x, event_dims)
def fake_log_prob(x):
"""Cooled version of the target distribution."""
return 1.1 * log_gamma_log_prob(x)
step_size = array_ops.placeholder(np.float32, [], name="step_size")
if feed_dict is None:
feed_dict = {}
feed_dict[step_size] = 0.4
sample, kernel_results = hmc.kernel(
target_log_prob_fn=log_gamma_log_prob,
current_state=initial_draws,
step_size=step_size,
num_leapfrog_steps=5,
seed=43)
bad_sample, bad_kernel_results = hmc.kernel(
target_log_prob_fn=fake_log_prob,
current_state=initial_draws,
step_size=step_size,
num_leapfrog_steps=5,
seed=44)
[
acceptance_probs_,
bad_acceptance_probs_,
initial_draws_,
updated_draws_,
fake_draws_,
] = sess.run([
kernel_results.acceptance_probs,
bad_kernel_results.acceptance_probs,
initial_draws,
sample,
bad_sample,
], feed_dict)
# Confirm step size is small enough that we usually accept.
self.assertGreater(acceptance_probs_.mean(), 0.5)
self.assertGreater(bad_acceptance_probs_.mean(), 0.5)
# Confirm step size is large enough that we sometimes reject.
self.assertLess(acceptance_probs_.mean(), 0.99)
self.assertLess(bad_acceptance_probs_.mean(), 0.99)
_, ks_p_value_true = stats.ks_2samp(initial_draws_.flatten(),
updated_draws_.flatten())
_, ks_p_value_fake = stats.ks_2samp(initial_draws_.flatten(),
fake_draws_.flatten())
logging_ops.vlog(
1, "acceptance rate for true target: {}".format(
acceptance_probs_.mean()))
logging_ops.vlog(
1, "acceptance rate for fake target: {}".format(
bad_acceptance_probs_.mean()))
logging_ops.vlog(
1, "K-S p-value for true target: {}".format(ks_p_value_true))
logging_ops.vlog(
1, "K-S p-value for fake target: {}".format(ks_p_value_fake))
# Make sure that the MCMC update hasn't changed the empirical CDF much.
self.assertGreater(ks_p_value_true, 1e-3)
# Confirm that targeting the wrong distribution does
# significantly change the empirical CDF.
self.assertLess(ks_p_value_fake, 1e-6)