code: def testReproduceBug159550941(self):
      # Reproduction for b/159550941.
      input_signature = [tf.TensorSpec([], tf.int32)]
  
      @tf.function(input_signature=input_signature)
      def sample(chains):
          initial_state = tf.zeros([chains, 1])
  
          def log_prob(x):
              return tf.reduce_sum(normal.Normal(0, 1).log_prob(x), -1)
  
          kernel = hmc.HamiltonianMonteCarlo(
              target_log_prob_fn=log_prob,
              num_leapfrog_steps=3,
              step_size=1e-3,
          )
          return sample_chain(
              num_results=5,
              num_burnin_steps=4,
              current_state=initial_state,
              kernel=kernel,
              trace_fn=None,
          )
  
      # Checking that shape inference doesn't fail.
      sample(2)
  
imports: import collections
  import warnings
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import pad
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  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 multinomial
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.mcmc import diagonal_mass_matrix_adaptation
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import nuts
  from tensorflow_probability.python.mcmc import random_walk_metropolis
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.sample import sample_chain
  from tensorflow_probability.python.random.random_ops import rayleigh
label_desc: Use image node_id 13 for calling the SampleChainTest obj's underlying member method code with example usage: obj.testReproduceBug159550941() and returns: sample_chain, tf