code: def test_reproduce_bug159550941(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,
          )
          results = sample_chain_with_burnin(
              num_results=5,
              num_burnin_steps=4,
              current_state=initial_state,
              kernel=kernel,
          )
          return results.trace
  
      # 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.distributions import normal
  from tensorflow_probability.python.experimental.mcmc import covariance_reducer
  from tensorflow_probability.python.experimental.mcmc import expectations_reducer
  from tensorflow_probability.python.experimental.mcmc.internal import test_fixtures
  from tensorflow_probability.python.experimental.mcmc.sample_fold import sample_chain_with_burnin
  from tensorflow_probability.python.experimental.mcmc.sample_fold import sample_fold
  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
label_desc: Use image node_id 7 for calling the SampleChainTest obj's underlying member method code with example usage: obj.test_reproduce_bug159550941() and returns: results, tf