code: def testCriterionStateEquivalence(self, criterion_fn):
      # Criteria should not care about the exact arrangement of state parts.
      previous_state = np.random.randn(4, 6).astype(self.dtype)
      new_state = np.random.randn(4, 6).astype(self.dtype)
      accept_prob = np.random.uniform(size=(4,)).astype(self.dtype)
  
      matrix_previous_state = previous_state.reshape([4, 3, 2])
      matrix_new_state = new_state.reshape([4, 3, 2])
  
      list_previous_state = [
          previous_state[:, :2],
          previous_state[:, 2:],
      ]
      list_new_state = [new_state[:, :2], new_state[:, 2:]]
  
      criterion = criterion_fn(
          previous_state, new_state, accept_prob, 1.0
      )
      matrix_criterion = criterion_fn(
          matrix_previous_state, matrix_new_state, accept_prob, 1.0
      )
      list_criterion = criterion_fn(
          list_previous_state, list_new_state, accept_prob, 1.0
      )
  
      self.assertAllEqual([4], criterion.shape)
      self.assertAllClose(criterion, matrix_criterion)
      self.assertAllClose(criterion, list_criterion)
  
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 half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.mcmc import gradient_based_trajectory_length_adaptation
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  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 dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import transformed_kernel
label_desc: Use image node_id 9 for calling the _GradientBasedTrajectoryLengthAdaptationTest obj's underlying member method code with example usage: obj.testCriterionStateEquivalence(criterion_fn) without return types