code: def test_gaussian_mixture_model(jit):
      K, N = 3, 1000
  
      def gmm(data):
          mix_proportions = pyro.sample(
              "phi", dist.Dirichlet(torch.ones(K))
          )
          with pyro.plate("num_clusters", K):
              cluster_means = pyro.sample(
                  "cluster_means",
                  dist.Normal(torch.arange(float(K)), 1.0),
              )
          with pyro.plate("data", data.shape[0]):
              assignments = pyro.sample(
                  "assignments", dist.Categorical(mix_proportions)
              )
              pyro.sample(
                  "obs",
                  dist.Normal(cluster_means[assignments], 1.0),
                  obs=data,
              )
          return cluster_means
  
      true_cluster_means = torch.tensor([1.0, 5.0, 10.0])
      true_mix_proportions = torch.tensor([0.1, 0.3, 0.6])
      cluster_assignments = dist.Categorical(
          true_mix_proportions
      ).sample(torch.Size((N,)))
      data = dist.Normal(
          true_cluster_means[cluster_assignments], 1.0
      ).sample()
      nuts_kernel = NUTS(
          gmm,
          max_plate_nesting=1,
          jit_compile=jit,
          ignore_jit_warnings=True,
      )
      mcmc = MCMC(nuts_kernel, num_samples=300, warmup_steps=100)
      mcmc.run(data)
      samples = mcmc.get_samples()
      assert_equal(
          samples["phi"].mean(0).sort()[0],
          true_mix_proportions,
          prec=0.05,
      )
      assert_equal(
          samples["cluster_means"].mean(0).sort()[0],
          true_cluster_means,
          prec=0.2,
      )
  
imports: import logging
  import os
  from collections import namedtuple
  import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.optim
  import pyro.poutine
  from pyro.contrib.conjugate.infer import BetaBinomialPair, GammaPoissonPair, collapse_conjugate, posterior_replay
  from pyro.infer import SVI, TraceEnum_ELBO
  from pyro.infer.autoguide import AutoDelta
  from pyro.infer.mcmc import MCMC, NUTS, ArrowheadMassMatrix
  from pyro.util import ignore_jit_warnings
  from tests.common import assert_close, assert_equal
  from .test_hmc import GaussianChain, rmse
label_desc: Use image node_id 9 for calling a global function with example usage: test_gaussian_mixture_model(jit) and returns: cluster_means