code: def testNested(self, bijector):
      if not tf.executing_eagerly():
          return
      base_mean = tf.constant([1.0, 0])
      base_cov = tf.constant([[1, 0.5], [0.5, 1]])
  
      dist_2d = tfd.MultivariateNormalTriL(
          loc=base_mean, scale_tril=tf.linalg.cholesky(base_cov)
      )
      dist_4d = tfd.MultivariateNormalDiag(scale_diag=tf.ones(4))
  
      target_dist = tfd.JointDistributionSequential(
          [
              bijector(dist_2d),
              tfb.Reshape([2, 2])(dist_4d),
          ]
      )
  
      kernel = neutra.NeuTra(
          target_log_prob_fn=lambda x, y: target_dist.log_prob((x, y)),
          state_shape=target_dist.event_shape,
          num_step_size_adaptation_steps=800,
          num_train_steps=1000,
          train_batch_size=64,
          seed=tfp_test_util.test_seed(),
          unconstraining_bijector=[bijector, tfb.Identity()],
      )
  
      chain_2d, chain_4d = tfp.mcmc.sample_chain(
          num_results=1000,
          num_burnin_steps=1000,
          current_state=tf.nest.map_structure(
              lambda s: tf.zeros([16] + s.as_list()),
              target_dist.event_shape,
          ),
          kernel=kernel,
          trace_fn=None,
          seed=tfp_test_util.test_seed(),
      )
  
      sample_mean_2d = tf.reduce_mean(
          input_tensor=chain_2d, axis=[0, 1]
      )
      sample_mean_4d = tf.reduce_mean(
          input_tensor=chain_4d, axis=[0, 1]
      )
  
      true_samples_2d, true_samples_4d = target_dist.sample(
          4096, seed=tfp_test_util.test_seed()
      )
  
      true_mean_2d = tf.reduce_mean(
          input_tensor=true_samples_2d, axis=0
      )
      true_mean_4d = tf.reduce_mean(
          input_tensor=true_samples_4d, axis=0
      )
  
      self.assertAllClose(
          true_mean_2d, sample_mean_2d, rtol=0.1, atol=0.1
      )
      self.assertAllClose(
          true_mean_4d, sample_mean_4d, 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 2 for calling the NeutraKernelTest obj's underlying member method code with example usage: obj.testNested(bijector) without return types