docstr: An end-to-end test of sampling using SMC.
docstr_tok: ["An","end-to-end","test","of","sampling","using","SMC","."]
code: 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.0,
              244.0,
              47.0,
              287.0,
              203.0,
              58.0,
              210.0,
              202.0,
              198.0,
              158.0,
              165.0,
              201.0,
              157.0,
              131.0,
              166.0,
              160.0,
              186.0,
              125.0,
              218.0,
              146.0,
          ]
      )
      obs = np.asarray(
          [
              592.0,
              401.0,
              583.0,
              402.0,
              495.0,
              173.0,
              479.0,
              504.0,
              510.0,
              416.0,
              393.0,
              442.0,
              317.0,
              311.0,
              400.0,
              337.0,
              423.0,
              334.0,
              533.0,
              344.0,
          ]
      )
      y_sigma = np.asarray(
          [
              61.0,
              25.0,
              38.0,
              15.0,
              21.0,
              15.0,
              27.0,
              14.0,
              30.0,
              16.0,
              14.0,
              25.0,
              52.0,
              16.0,
              34.0,
              31.0,
              42.0,
              26.0,
              16.0,
              22.0,
          ]
      )
      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 = jds.JointDistributionSequential(
          [
              normal.Normal(loc=hyper_mean, scale=hyper_scale),
              normal.Normal(loc=hyper_mean, scale=hyper_scale),
              normal.Normal(loc=hyper_mean, scale=hyper_scale),
              half_normal.HalfNormal(scale=tf.cast(0.5, dtype)),
              uniform.Uniform(low=tf.cast(0, dtype), high=0.5),
          ],
          validate_args=True,
      )
  
      def likelihood_log_prob_fn(b0, b1, mu_out, sigma_out, weight):
          return independent.Independent(
              mixture.Mixture(
                  categorical.Categorical(
                      probs=tf.stack(
                          [
                              tf.repeat(
                                  1 - weight[..., tf.newaxis],
                                  20,
                                  axis=-1,
                              ),
                              tf.repeat(
                                  weight[..., tf.newaxis], 20, axis=-1
                              ),
                          ],
                          -1,
                      )
                  ),
                  [
                      normal.Normal(
                          loc=b0[..., tf.newaxis]
                          + b1[..., tf.newaxis] * predictors,
                          scale=y_sigma,
                      ),
                      normal.Normal(
                          loc=mu_out[..., tf.newaxis],
                          scale=y_sigma + sigma_out[..., tf.newaxis],
                      ),
                  ],
              ),
              1,
          ).log_prob(obs)
  
      unconstraining_bijectors = [
          identity.Identity(),
          identity.Identity(),
          identity.Identity(),
          softplus.Softplus(),
          sigmoid.Sigmoid(tf.constant(0.0, dtype), 0.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 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=0.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
      )
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mixture
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.distributions.internal import statistical_testing
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import compute_hmc_step_size
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import gen_make_hmc_kernel_fn
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import gen_make_transform_hmc_kernel_fn
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import make_rwmh_kernel_fn
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import sample_sequential_monte_carlo
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import simple_heuristic_tuning
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import nuts
label_desc: Use image node_id 4 for calling the SampleSequentialMonteCarloTest obj's underlying member method code with example usage: obj.testSampleEndtoEndXLA() and returns: independent, sample_sequential_monte_carlo