code: def testLatentsOfMixedRank(self, batch_shape, num_steps):
      strm = test_util.test_seed_stream()
  
      init0 = [tf.ones(batch_shape + [6])]
      init1 = [
          tf.ones(batch_shape + []),
          tf.ones(batch_shape + [1]),
          tf.ones(batch_shape + [2, 2]),
      ]
  
      @tf.function(autograph=False)
      def run_two_chains(init0, init1):
          def log_prob0(x):
              return tf.squeeze(
                  independent.Independent(
                      normal.Normal(
                          tf.range(6, dtype=tf.float32),
                          tf.constant(1.0),
                      ),
                      reinterpreted_batch_ndims=1,
                  ).log_prob(x)
              )
  
          kernel0 = nuts.NoUTurnSampler(log_prob0, step_size=0.3)
          [results0] = sample.sample_chain(
              num_results=num_steps,
              num_burnin_steps=10,
              current_state=init0,
              kernel=kernel0,
              trace_fn=None,
              seed=strm(),
          )
  
          def log_prob1(state0, state1, state2):
              return tf.squeeze(
                  normal.Normal(
                      tf.constant(0.0), tf.constant(1.0)
                  ).log_prob(state0)
                  + independent.Independent(
                      normal.Normal(
                          tf.constant([1.0]), tf.constant(1.0)
                      ),
                      reinterpreted_batch_ndims=1,
                  ).log_prob(state1)
                  + independent.Independent(
                      normal.Normal(
                          tf.constant([[2.0, 3.0], [4.0, 5.0]]),
                          tf.constant(1.0),
                      ),
                      reinterpreted_batch_ndims=2,
                  ).log_prob(state2)
              )
  
          kernel1 = nuts.NoUTurnSampler(log_prob1, step_size=0.3)
          results1_ = sample.sample_chain(
              num_results=num_steps,
              num_burnin_steps=10,
              current_state=init1,
              kernel=kernel1,
              trace_fn=None,
              seed=strm(),
          )
          results1 = tf.concat(
              [
                  tf.reshape(x, [num_steps] + batch_shape + [-1])
                  for x in results1_
              ],
              axis=-1,
          )
  
          return results0, results1
  
      results0, results1 = run_two_chains(init0, init1)
  
      self.evaluate(
          st.assert_true_cdf_equal_by_dkwm_two_sample(
              results0, results1
          )
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import sigmoid_beta
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions.internal import statistical_testing
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import nuts
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 8 for calling the NutsTest obj's underlying member method code with example usage: obj.testLatentsOfMixedRank(batch_shape, num_steps) and returns: results0,  results1, tf, tf