code: def test_multistate(self):
      n_samples = 5
      # state_0 is two scalar chains taken from iid Normal(0, 1).
      state_0 = np.random.randn(n_samples, 2)
  
      # state_1 is three 4-variate chains taken from Normal(0, 1) that have been
      # shifted.
      offset = np.array([1.0, -1.0, 2.0]).reshape(3, 1)
      state_1 = np.random.randn(n_samples, 3, 4) + offset
      running_rhat = (
          sample_stats.RunningPotentialScaleReduction.from_shape(
              shape=[(2,), (3, 4)], independent_chain_ndims=[1, 1]
          )
      )
      for sample in zip(state_0, state_1):
          running_rhat = running_rhat.update(sample)
      rhat = self.evaluate(running_rhat.potential_scale_reduction())
      true_rhat = diagnostic.potential_scale_reduction(
          chains_states=[state_0, state_1], independent_chain_ndims=1
      )
      self.assertAllCloseNested(true_rhat, rhat, rtol=1e-6)
  
imports: from absl.testing import parameterized
  import numpy
  import scipy.stats
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import diagnostic
label_desc: Use image node_id 3 for calling the RunningPotentialScaleReductionTest obj's underlying member method code with example usage: obj.test_multistate() without return types