docstr: An end-to-end test of sampling using NUTS.
code: def testSampleEndtoEndXLA(self):
      """An end-to-end test of sampling using NUTS."""
      if (
          tf.executing_eagerly()
          or tf.config.experimental_functions_run_eagerly()
      ):
          self.skipTest(
              "No need to test XLA under all execution regimes."
          )
  
      strm = test_util.test_seed_stream()
      predictors = tf.cast(
          [
              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,
          ],
          tf.float32,
      )
      obs = tf.cast(
          [
              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,
          ],
          tf.float32,
      )
      y_sigma = tf.cast(
          [
              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,
          ],
          tf.float32,
      )
  
      # Robust linear regression model
      robust_lm = jds.JointDistributionSequential(
          [
              normal.Normal(loc=0.0, scale=1.0),  # b0
              normal.Normal(loc=0.0, scale=1.0),  # b1
              half_normal.HalfNormal(5.0),  # df
              lambda df, b1, b0: independent.Independent(  # pylint: disable=g-long-lambda
                  student_t.StudentT(  # Likelihood
                      df=df[..., tf.newaxis],
                      loc=(
                          b0[..., tf.newaxis]
                          + b1[..., tf.newaxis] * predictors[tf.newaxis]
                      ),
                      scale=y_sigma,
                  )
              ),
          ],
          validate_args=True,
      )
  
      log_prob = lambda b0, b1, df: robust_lm.log_prob(
          [b0, b1, df, obs]
      )
      init_step_size = [1.0, 0.2, 0.5]
      step_size0 = [
          tf.cast(x, dtype=tf.float32) for x in init_step_size
      ]
  
      number_of_steps, burnin, nchain = 200, 50, 10
  
      @tf.function(autograph=False, jit_compile=True)
      def run_chain_and_get_diagnostic():
          # Ensure we're really in graph mode.
          assert hasattr(tf.constant([]), "graph")
  
          # random initialization of the starting postion of each chain
          b0, b1, df, _ = robust_lm.sample(nchain, seed=strm())
  
          # bijector to map contrained parameters to real
          unconstraining_bijectors = [
              identity.Identity(),
              identity.Identity(),
              exp.Exp(),
          ]
  
          def trace_fn(_, pkr):
              return (
                  pkr.inner_results.inner_results.step_size,
                  pkr.inner_results.inner_results.log_accept_ratio,
              )
  
          kernel = duassa.DualAveragingStepSizeAdaptation(
              transformed_kernel.TransformedTransitionKernel(
                  inner_kernel=nuts.NoUTurnSampler(
                      target_log_prob_fn=log_prob, step_size=step_size0
                  ),
                  bijector=unconstraining_bijectors,
              ),
              target_accept_prob=0.8,
              num_adaptation_steps=burnin,
          )
  
          # Sampling from the chain and get diagnostics
          mcmc_trace, (
              step_size,
              log_accept_ratio,
          ) = sample.sample_chain(
              num_results=number_of_steps,
              num_burnin_steps=burnin,
              current_state=[b0, b1, df],
              kernel=kernel,
              trace_fn=trace_fn,
              seed=strm(),
          )
          rhat = diagnostic.potential_scale_reduction(mcmc_trace)
          return (
              [s[-1] for s in step_size],  # final step size
              tf.math.exp(generic.reduce_logmeanexp(log_accept_ratio)),
              [tf.reduce_mean(rhat_) for rhat_ in rhat],  # average rhat
          )
  
      # Sample from posterior distribution and get diagnostic
      [
          final_step_size,
          average_accept_ratio,
          average_rhat,
      ] = self.evaluate(run_chain_and_get_diagnostic())
  
      # Check that step size adaptation reduced the initial step size
      self.assertAllLess(
          np.asarray(final_step_size) - np.asarray(init_step_size), 0.0
      )
      # Check that average acceptance ratio is close to target
      self.assertAllClose(
          average_accept_ratio,
          0.8 * np.ones_like(average_accept_ratio),
          atol=0.1,
          rtol=0.1,
      )
      # Check that mcmc sample quality is acceptable with tuning
      self.assertAllClose(
          average_rhat, np.ones_like(average_rhat), atol=0.05, rtol=0.05
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import sigmoid_beta
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions.internal import statistical_testing
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import nuts
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 13 for calling the NutsTest obj's underlying member method code with example usage: obj.testSampleEndtoEndXLA() and returns: unknown,  tf,  unknown, pkr,  pkr