code: def test_gaussian_hmm(num_steps):
      dim = 4
  
      def model(data):
          initialize = pyro.sample(
              "initialize", dist.Dirichlet(torch.ones(dim))
          )
          with pyro.plate("states", dim):
              transition = pyro.sample(
                  "transition", dist.Dirichlet(torch.ones(dim, dim))
              )
              emission_loc = pyro.sample(
                  "emission_loc",
                  dist.Normal(torch.zeros(dim), torch.ones(dim)),
              )
              emission_scale = pyro.sample(
                  "emission_scale",
                  dist.LogNormal(torch.zeros(dim), torch.ones(dim)),
              )
          x = None
          with ignore_jit_warnings(
              [("Iterating over a tensor", RuntimeWarning)]
          ):
              for t, y in pyro.markov(enumerate(data)):
                  x = pyro.sample(
                      "x_{}".format(t),
                      dist.Categorical(
                          initialize if x is None else transition[x]
                      ),
                      infer={"enumerate": "parallel"},
                  )
                  pyro.sample(
                      "y_{}".format(t),
                      dist.Normal(emission_loc[x], emission_scale[x]),
                      obs=y,
                  )
  
      def _get_initial_trace():
          guide = AutoDelta(
              poutine.block(
                  model,
                  expose_fn=lambda msg: not msg["name"].startswith("x")
                  and not msg["name"].startswith("y"),
              )
          )
          elbo = TraceEnum_ELBO(max_plate_nesting=1)
          svi = SVI(model, guide, optim.Adam({"lr": 0.01}), elbo)
          for _ in range(100):
              svi.step(data)
          return poutine.trace(guide).get_trace(data)
  
      def _generate_data():
          transition_probs = torch.rand(dim, dim)
          emissions_loc = torch.arange(dim, dtype=torch.Tensor().dtype)
          emissions_scale = 1.0
          state = torch.tensor(1)
          obs = [
              dist.Normal(
                  emissions_loc[state], emissions_scale
              ).sample()
          ]
          for _ in range(num_steps):
              state = dist.Categorical(transition_probs[state]).sample()
              obs.append(
                  dist.Normal(
                      emissions_loc[state], emissions_scale
                  ).sample()
              )
          return torch.stack(obs)
  
      data = _generate_data()
      nuts_kernel = NUTS(
          model,
          max_plate_nesting=1,
          jit_compile=True,
          ignore_jit_warnings=True,
      )
      if num_steps == 30:
          nuts_kernel.initial_trace = _get_initial_trace()
      mcmc = MCMC(nuts_kernel, num_samples=5, warmup_steps=5)
      mcmc.run(data)
  
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 11 for calling a global function with example usage: test_gaussian_hmm(num_steps) and returns: poutine, torch