code: def _tfTestHelper(self, run_asserts_fn, execute_program_fn):
      # Note: test_device is 'cpu', 'gpu', etc.
  
      # Various int32 and int64 kernels are missing for GPU, so we skip direct
      # tests on the GPU device, but test XLA on GPU below.
      if "cpu" in FLAGS.test_device.lower():
          # Make sure everything works with no XLA compilation.
          with tf.device("CPU:0"):
              run_asserts_fn(
                  functools.partial(
                      execute_program_fn, backend=TF_BACKEND
                  )
              )
  
      # Force XLA compilation using tf.function.
      backend = TF_BACKEND_NO_ASSERTS
      f = functools.partial(execute_program_fn, backend=backend)
      f = tf.function(f, autograph=False, jit_compile=True)
      with tf.device(FLAGS.test_device):
          run_asserts_fn(f)
  
imports: import functools
  from absl import flags
  import numpy
  import tensorflow.compat.v1
  from tensorflow_probability.python.experimental.auto_batching import numpy_backend
  from tensorflow_probability.python.experimental.auto_batching import test_programs
  from tensorflow_probability.python.experimental.auto_batching import tf_backend
  from tensorflow_probability.python.experimental.auto_batching import virtual_machine
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the VMTest obj's underlying member method code with example usage: obj._tfTestHelper(run_asserts_fn, execute_program_fn) without return types