code: def test_mcmc_interface(
      run_mcmc_cls, num_draws, group_by_chain, num_chains
  ):
      num_samples = 2000
      data = torch.tensor([1.0])
      initial_params, _, transforms, _ = initialize_model(
          normal_normal_model, model_args=(data,), num_chains=num_chains
      )
      kernel = PriorKernel(normal_normal_model)
      samples, mcmc_num_chains = run_mcmc_cls(
          data,
          kernel,
          num_samples=num_samples,
          warmup_steps=100,
          initial_params=initial_params,
          num_chains=num_chains,
          mp_context="spawn",
          transforms=transforms,
          num_draws=num_draws,
          group_by_chain=group_by_chain,
      )
      # test sample shape
      expected_samples = (
          num_draws if num_draws is not None else num_samples
      )
      if group_by_chain:
          expected_shape = (mcmc_num_chains, expected_samples, 1)
      elif num_draws is not None:
          # FIXME: what is the expected behavior of num_draw is not None and group_by_chain=False?
          expected_shape = (expected_samples, 1)
      else:
          expected_shape = (mcmc_num_chains * expected_samples, 1)
      assert samples["y"].shape == expected_shape
  
      # test sample stats
      if group_by_chain:
          samples = {
              k: v.reshape((-1,) + v.shape[2:])
              for k, v in samples.items()
          }
      sample_mean = samples["y"].mean()
      sample_std = samples["y"].std()
      assert_close(sample_mean, torch.tensor(0.0), atol=0.1)
      assert_close(sample_std, torch.tensor(1.0), atol=0.1)
  
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 4 for calling a global function with example usage: test_mcmc_interface(run_mcmc_cls, num_draws, group_by_chain, num_chains) without return types