code: def test_benchmark_sampling_with_xla(self):
      if not tf.executing_eagerly():
          return
      seed = test_util.test_seed()
  
      @tf.function(autograph=False, jit_compile=True)
      def _run():
          (
              model,
              observed_time_series,
              is_missing,
          ) = self._build_test_model(num_timesteps=336, batch_shape=[])
          return gibbs_sampler.fit_with_gibbs_sampling(
              model,
              missing_values_util.MaskedTimeSeries(
                  observed_time_series[..., tf.newaxis], is_missing
              ),
              num_results=500,
              num_warmup_steps=100,
              seed=seed,
          )
  
      t0 = time.time()
      samples = _run()
      t1 = time.time()
      print(
          "Drew (100+500) samples (with JIT compilation) in time",
          t1 - t0,
      )
  
      t0 = time.time()
      samples = _run()
      t1 = time.time()
      print("Drew (100+500) samples in time", t1 - t0)
  
      print("Results:", samples)
  
file_path: probability/tensorflow_probability/python/experimental/sts_gibbs/benchmarks_test.py
imports: import time
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.experimental.sts_gibbs import gibbs_sampler
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts.internal import missing_values_util
label_desc: Use image node_id 2 for calling the XLABenchmarkTests obj's underlying member method code with example usage: obj.test_benchmark_sampling_with_xla() and returns: gibbs_sampler