def testCorrectStepSizeTransformedkernel(self): scalings = .1 bijector = tfb.Sigmoid() prior = tfd.Beta(.1, .1) likelihood = tfd.Beta(5., 5.) init_state = [tf.clip_by_value(prior.sample(10000), 1e-5, 1. - 1e-5)] make_transform_kernel_fn = gen_make_transform_hmc_kernel_fn( [bijector], num_leapfrog_steps=1) kernel = make_transform_kernel_fn(likelihood.log_prob, init_state, scalings=scalings) step_size, expected_step_size = self.evaluate([ tf.squeeze(kernel.inner_kernel.step_size), scalings * tf.math.reduce_std(bijector.inverse(init_state)) ]) self.assertAllGreater(step_size, 0.) self.assertAllEqual(step_size, expected_step_size)
def testSampleEndtoEndXLA(self): """An end-to-end test of sampling using SMC.""" if tf.executing_eagerly( ) or tf.config.experimental_functions_run_eagerly(): self.skipTest('No need to test XLA under all execution regimes.') seed = test_util.test_seed() dtype = tf.float32 # Set up data. predictors = np.asarray([ 201., 244., 47., 287., 203., 58., 210., 202., 198., 158., 165., 201., 157., 131., 166., 160., 186., 125., 218., 146. ]) obs = np.asarray([ 592., 401., 583., 402., 495., 173., 479., 504., 510., 416., 393., 442., 317., 311., 400., 337., 423., 334., 533., 344. ]) y_sigma = np.asarray([ 61., 25., 38., 15., 21., 15., 27., 14., 30., 16., 14., 25., 52., 16., 34., 31., 42., 26., 16., 22. ]) y_sigma = tf.cast(y_sigma / (2 * obs.std(axis=0)), dtype) obs = tf.cast((obs - obs.mean(axis=0)) / (2 * obs.std(axis=0)), dtype) predictors = tf.cast((predictors - predictors.mean(axis=0)) / (2 * predictors.std(axis=0)), dtype) hyper_mean = tf.cast(0, dtype) hyper_scale = tf.cast(2.5, dtype) # Generate model prior_log_prob_fn and likelihood_log_prob_fn. prior_jd = tfd.JointDistributionSequential([ tfd.Normal(loc=hyper_mean, scale=hyper_scale), tfd.Normal(loc=hyper_mean, scale=hyper_scale), tfd.Normal(loc=hyper_mean, scale=hyper_scale), tfd.HalfNormal(scale=tf.cast(.5, dtype)), tfd.Uniform(low=tf.cast(0, dtype), high=.5), ], validate_args=True) def likelihood_log_prob_fn(b0, b1, mu_out, sigma_out, weight): return tfd.Independent( tfd.Mixture( tfd.Categorical(probs=tf.stack([ tf.repeat(1 - weight[..., tf.newaxis], 20, axis=-1), tf.repeat(weight[..., tf.newaxis], 20, axis=-1) ], -1)), [ tfd.Normal(loc=b0[..., tf.newaxis] + b1[..., tf.newaxis] * predictors, scale=y_sigma), tfd.Normal(loc=mu_out[..., tf.newaxis], scale=y_sigma + sigma_out[..., tf.newaxis]) ]), 1).log_prob(obs) unconstraining_bijectors = [ tfb.Identity(), tfb.Identity(), tfb.Identity(), tfb.Softplus(), tfb.Sigmoid(tf.constant(0., dtype), .5), ] make_transform_hmc_kernel_fn = gen_make_transform_hmc_kernel_fn( unconstraining_bijectors, num_leapfrog_steps=5) @tf.function(autograph=False, jit_compile=True) def run_smc(): # Ensure we're really in graph mode. assert hasattr(tf.constant([]), 'graph') return tfp.experimental.mcmc.sample_sequential_monte_carlo( prior_jd.log_prob, likelihood_log_prob_fn, prior_jd.sample([1000, 5], seed=seed), make_kernel_fn=make_transform_hmc_kernel_fn, tuning_fn=functools.partial(simple_heuristic_tuning, optimal_accept=.6), min_num_steps=5, seed=seed) n_stage, (b0, b1, mu_out, sigma_out, weight), _ = run_smc() (n_stage, b0, b1, mu_out, sigma_out, weight) = self.evaluate( (n_stage, b0, b1, mu_out, sigma_out, weight)) self.assertTrue(n_stage, 10) # Compare the SMC posterior with the result from a calibrated HMC. self.assertAllClose(tf.reduce_mean(b0), 0.016, atol=0.005, rtol=0.005) self.assertAllClose(tf.reduce_mean(b1), 1.245, atol=0.005, rtol=0.035) self.assertAllClose(tf.reduce_mean(weight), 0.28, atol=0.03, rtol=0.02) self.assertAllClose(tf.reduce_mean(mu_out), 0.13, atol=0.2, rtol=0.2) self.assertAllClose(tf.reduce_mean(sigma_out), 0.46, atol=0.5, rtol=0.5) self.assertAllClose(tf.math.reduce_std(b0), 0.031, atol=0.015, rtol=0.3) self.assertAllClose(tf.math.reduce_std(b1), 0.068, atol=0.1, rtol=0.1) self.assertAllClose(tf.math.reduce_std(weight), 0.1, atol=0.1, rtol=0.1)