code: def testSingleTensor(self, bijector, dtype):
      if not tf.executing_eagerly():
          return
      base_mean = tf.convert_to_tensor(value=[1.0, 0], dtype=dtype)
      base_cov = tf.convert_to_tensor(
          value=[[1, 0.5], [0.5, 1]], dtype=dtype
      )
  
      base_dist = tfd.MultivariateNormalTriL(
          loc=base_mean, scale_tril=tf.linalg.cholesky(base_cov)
      )
      target_dist = bijector(base_dist)
  
      def debug_fn(*args):
          del args
          debug_fn.count += 1
  
      debug_fn.count = 0
  
      kernel = neutra.NeuTra(
          target_log_prob_fn=target_dist.log_prob,
          state_shape=2,
          num_step_size_adaptation_steps=800,
          num_train_steps=1000,
          train_batch_size=64,
          learning_rate=tf.convert_to_tensor(value=1e-2, dtype=dtype),
          seed=tfp_test_util.test_seed(),
          train_debug_fn=debug_fn,
          unconstraining_bijector=bijector,
      )
  
      chain = tfp.mcmc.sample_chain(
          num_results=1000,
          num_burnin_steps=1000,
          current_state=tf.zeros([16, 2], dtype),
          kernel=kernel,
          trace_fn=None,
          seed=tfp_test_util.test_seed(),
      )
      self.assertEqual(1000, debug_fn.count)
  
      sample_mean = tf.reduce_mean(input_tensor=chain, axis=[0, 1])
      sample_cov = tfp.stats.covariance(chain, sample_axis=[0, 1])
  
      true_samples = target_dist.sample(
          4096, seed=tfp_test_util.test_seed()
      )
  
      true_mean = tf.reduce_mean(input_tensor=true_samples, axis=0)
      true_cov = tfp.stats.covariance(chain, sample_axis=[0, 1])
  
      self.assertAllClose(true_mean, sample_mean, rtol=0.1, atol=0.1)
      self.assertAllClose(true_cov, sample_cov, rtol=0.1, atol=0.1)
  
file_path: probability/discussion/neutra/neutra_kernel_test.py
imports: from absl.testing import parameterized
  import tensorflow.compat.v1
  import tensorflow_probability
  from discussion import neutra
  from tensorflow_probability.python.internal import test_util
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 1 for calling the NeutraKernelTest obj's underlying member method code with example usage: obj.testSingleTensor(bijector, dtype) without return types