code: def test_jit_works_correctly(self):
      @tf.function(jit_compile=True)
      def sample_float32():
          return sample_halton_sequence_lib.sample_halton_sequence(
              5,
              num_results=10,
              dtype=tf.float32,
              seed=test_util.test_seed(),
          )
  
      samples = sample_float32()
      self.assertEqual(samples.shape, [10, 5])
  
      @tf.function(jit_compile=True)
      def sample_float64():
          return sample_halton_sequence_lib.sample_halton_sequence(
              5,
              num_results=10,
              dtype=tf.float64,
              seed=test_util.test_seed(),
          )
  
      samples = sample_float64()
      self.assertEqual(samples.shape, [10, 5])
  
file_path: probability/tensorflow_probability/python/mcmc/sample_halton_sequence_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import monte_carlo
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import sample_halton_sequence_lib
label_desc: Use image node_id 5 for calling the HaltonSequenceTest obj's underlying member method code with example usage: obj.test_jit_works_correctly() and returns: sample_halton_sequence_lib, sample_halton_sequence_lib