code: def test_infer_discrete(temperature, length):
      @ignore_jit_warnings()
      def hmm(transition, means, data):
          states = [torch.tensor(0)]
          for t in pyro.markov(range(len(data))):
              states.append(
                  pyro.sample(
                      "states_{}".format(t),
                      dist.Categorical(transition[states[-1]]),
                      infer={"enumerate": "parallel"},
                  )
              )
              pyro.sample(
                  "obs_{}".format(t),
                  dist.Normal(means[states[-1]], 1.0),
                  obs=data[t],
              )
          return tuple(states)
  
      hidden_dim = 10
      transition = 0.3 / hidden_dim + 0.7 * torch.eye(hidden_dim)
      means = torch.arange(float(hidden_dim))
      data = 1 + 2 * torch.randn(length)
  
      decoder = infer_discrete(
          hmm, first_available_dim=-1, temperature=temperature
      )
      jit_decoder = pyro.ops.jit.trace(decoder)
  
      states = decoder(transition, means, data)
      jit_states = jit_decoder(transition, means, data)
      assert len(states) == len(jit_states)
      for state, jit_state in zip(states, jit_states):
          assert state.shape == jit_state.shape
          if temperature == 0:
              assert_equal(state, jit_state)
  
imports: import logging
  import warnings
  import pytest
  import torch
  from torch.autograd import grad
  from torch.distributions import constraints, kl_divergence
  import pyro
  import pyro.distributions
  import pyro.ops.jit
  import pyro.poutine
  from pyro.distributions.util import scale_and_mask
  from pyro.infer import SVI, JitTrace_ELBO, JitTraceEnum_ELBO, JitTraceGraph_ELBO, JitTraceMeanField_ELBO, Trace_ELBO, TraceEnum_ELBO, TraceGraph_ELBO, TraceMeanField_ELBO, TraceTMC_ELBO, infer_discrete
  from pyro.optim import Adam
  from pyro.poutine.indep_messenger import CondIndepStackFrame
  from pyro.util import ignore_jit_warnings
  from tests.common import assert_close, assert_equal
label_desc: Use image node_id 21 for calling a global function with example usage: test_infer_discrete(temperature, length) and returns: tuple