code: def test_batches(self):
      mvn = jds.JointDistributionSequential(
          [
              normal.Normal(1.0, 0.1),
              normal.Normal(2.0, 1.0),
              normal.Normal(3.0, 10.0),
          ]
      )
      n_chains = 10
      if self.use_default_momentum_distribution:
          momentum_distribution = None
          step_size = 0.1
      else:
          reshape_to_scalar = reshape.Reshape(event_shape_out=[])
          momentum_distribution = jds.JointDistributionSequential(
              [
                  reshape_to_scalar(
                      mvnpflo.MultivariateNormalPrecisionFactorLinearOperator(
                          precision_factor=tf.linalg.LinearOperatorDiag(
                              tf.fill([n_chains, 1], 0.1)
                          )
                      )
                  ),
                  reshape_to_scalar(
                      mvnpflo.MultivariateNormalPrecisionFactorLinearOperator(
                          precision_factor=tf.linalg.LinearOperatorDiag(
                              tf.fill([n_chains, 1], 1.0)
                          )
                      )
                  ),
                  reshape_to_scalar(
                      mvnpflo.MultivariateNormalPrecisionFactorLinearOperator(
                          precision_factor=tf.linalg.LinearOperatorDiag(
                              tf.fill([n_chains, 1], 10.0)
                          )
                      )
                  ),
              ]
          )
          step_size = 0.3
  
      hmc_kernel = phmc.PreconditionedHamiltonianMonteCarlo(
          target_log_prob_fn=mvn.log_prob,
          momentum_distribution=momentum_distribution,
          step_size=step_size,
          num_leapfrog_steps=10,
      )
  
      draws = sample.sample_chain(
          100,
          [tf.zeros([n_chains]) for _ in range(3)],
          kernel=hmc_kernel,
          seed=test_util.test_seed(),
          trace_fn=None,
      )
      ess = diagnostic.effective_sample_size(
          draws,
          cross_chain_dims=[1 for _ in draws],
          filter_threshold=0,
          filter_beyond_positive_pairs=False,
      )
      if not self.use_default_momentum_distribution:
          self.assertAllClose(
              self.evaluate(ess), 100 * n_chains * tf.ones(3)
          )
      else:
          self.assertLess(self.evaluate(tf.reduce_min(ess)), 100.0)
  
imports: import collections
  from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import wishart
  from tensorflow_probability.python.experimental.distributions import mvn_precision_factor_linop
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.internal import auto_composite_tensor
  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.internal import unnest
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 9 for calling the _PreconditionedHMCTest obj's underlying member method code with example usage: obj.test_batches() without return types