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)