code: def test_nuts_samples_well(self):
      if tf.executing_eagerly() and not JAX_MODE:
          self.skipTest("Eager test is slow.")
      model = eight_schools_named()
      pins = {"treatment_effects": tf.constant(TREATMENT_EFFECTS)}
  
      @tf.function
      def do_sample():
          return windowed_sampling.windowed_adaptive_nuts(
              200,
              model,
              max_tree_depth=5,
              seed=test_util.test_seed(),
              **pins,
          )
  
      draws, _ = do_sample()
      flat_draws = tf.nest.flatten(
          model.experimental_pin(**pins)._model_flatten(draws)
      )
      max_scale_reduction = tf.reduce_max(
          tf.nest.map_structure(
              tf.reduce_max,
              diagnostic.potential_scale_reduction(flat_draws),
          )
      )
      self.assertLess(self.evaluate(max_scale_reduction), 1.05)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import pad
  from tensorflow_probability.python.distributions import autoregressive
  from tensorflow_probability.python.distributions import beta_binomial
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import multinomial
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental import distribute
  from tensorflow_probability.python.experimental.mcmc import windowed_sampling
  from tensorflow_probability.python.internal import callable_util
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.mcmc import diagnostic
label_desc: Use image node_id 4 for calling the WindowedSamplingTest obj's underlying member method code with example usage: obj.test_nuts_samples_well() and returns: windowed_sampling