code: def test_multiple_latent_states_and_independent_chain_ndims(self):
      rng = test_util.test_np_rng()
      rhat_reducer = psrr.PotentialScaleReductionReducer(
          independent_chain_ndims=2
      )
      state = rhat_reducer.initialize(
          [tf.zeros((2, 5, 3)), tf.zeros((7, 2, 8))]
      )
      chain_state = rng.randn(4, 2, 5, 3)
      second_chain_state = rng.randn(4, 7, 2, 8)
      for latent in zip(chain_state, second_chain_state):
          state = rhat_reducer.one_step(latent, state)
      rhat = rhat_reducer.finalize(state)
      true_rhat = diagnostic.potential_scale_reduction(
          chains_states=[chain_state, second_chain_state],
          independent_chain_ndims=2,
      )
      rhat, true_rhat = self.evaluate([rhat, true_rhat])
      self.assertAllClose(true_rhat, rhat, rtol=1e-6)
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/potential_scale_reduction_reducer_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.experimental.mcmc import potential_scale_reduction_reducer
  from tensorflow_probability.python.experimental.mcmc import tracing_reducer
  from tensorflow_probability.python.experimental.mcmc.internal import test_fixtures
  from tensorflow_probability.python.experimental.mcmc.sample_fold import sample_fold
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow_probability.python.mcmc import hmc
label_desc: Use image node_id 5 for calling the PotentialScaleReductionReducerTest obj's underlying member method code with example usage: obj.test_multiple_latent_states_and_independent_chain_ndims() without return types