code: def testSampleXLA(self):
      self.skip_if_no_xla()
      if not tf.executing_eagerly():
          return  # jit_compile is eager-only.
  
      try:
          old_random_poisson = poisson_lib.random_poisson
          # TODO(b/195975508): Reloading the function to reset the cache.
          if not test_util.JAX_MODE:
              poisson_lib.random_poisson = tf.function(
                  poisson_lib.random_poisson._python_function
              )
  
          log_rates = np.random.rand(4, 3).astype(np.float32)
          dist = poisson_lib.Poisson(
              log_rate=log_rates, validate_args=True
          )
          # Verify the compile succeeds going all the way through the distribution.
          self.evaluate(
              tf.function(
                  lambda: dist.sample(5, seed=test_util.test_seed()),
                  jit_compile=True,
              )()
          )
          # Also test the low-level sampler and verify the XLA-friendly variant.
          _, runtime = self.evaluate(
              tf.function(poisson_lib.random_poisson, jit_compile=True)(
                  shape=[],
                  rates=tf.constant(10.0),
                  seed=test_util.test_seed(),
              )
          )
          self.assertEqual(
              implementation_selection._RUNTIME_DEFAULT, runtime
          )
      finally:
          poisson_lib.random_poisson = old_random_poisson
  
imports: import numpy
  from scipy import stats
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions.internal import statistical_testing
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import implementation_selection
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 3 for calling the PoissonSamplingTest obj's underlying member method code with example usage: obj.testSampleXLA() without return types