code: def test_mcmc_diagnostics(run_mcmc_cls, num_chains):
      data = torch.tensor([2.0]).repeat(3)
      initial_params, _, transforms, _ = initialize_model(
          normal_normal_model, model_args=(data,), num_chains=num_chains
      )
      kernel = PriorKernel(normal_normal_model)
      if run_mcmc_cls == run_default_mcmc:
          mcmc = MCMC(
              kernel,
              num_samples=10,
              warmup_steps=10,
              num_chains=num_chains,
              mp_context="spawn",
              initial_params=initial_params,
              transforms=transforms,
          )
      else:
          mcmc = StreamingMCMC(
              kernel,
              num_samples=10,
              warmup_steps=10,
              num_chains=num_chains,
              initial_params=initial_params,
              transforms=transforms,
          )
      mcmc.run(data)
      if not torch.backends.mkl.is_available():
          pytest.skip()
      diagnostics = mcmc.diagnostics()
      if (
          run_mcmc_cls == run_default_mcmc
      ):  # TODO n_eff for streaming MCMC
          assert diagnostics["y"]["n_eff"].shape == data.shape
      assert diagnostics["y"]["r_hat"].shape == data.shape
      assert diagnostics["dummy_key"] == {
          "chain {}".format(i): "dummy_value" for i in range(num_chains)
      }
  
imports: import os
  from functools import partial
  import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.infer.mcmc import HMC, NUTS
  from pyro.infer.mcmc.api import MCMC, StreamingMCMC, _MultiSampler, _UnarySampler
  from pyro.infer.mcmc.mcmc_kernel import MCMCKernel
  from pyro.infer.mcmc.util import initialize_model, select_samples
  from pyro.ops.streaming import StackStats, StatsOfDict
  from pyro.util import optional
  from tests.common import assert_close
label_desc: Use image node_id 9 for calling a global function with example usage: test_mcmc_diagnostics(run_mcmc_cls, num_chains) without return types