code: def test_traceenum_elbo(length):
      hidden_dim = 10
      transition = pyro.param(
          "transition",
          0.3 / hidden_dim + 0.7 * torch.eye(hidden_dim),
          constraint=constraints.positive,
      )
      means = pyro.param("means", torch.arange(float(hidden_dim)))
      data = 1 + 2 * torch.randn(length)
  
      @ignore_jit_warnings()
      def model(data):
          transition = pyro.param("transition")
          means = pyro.param("means")
          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)
  
      def guide(data):
          pass
  
      expected_loss = TraceEnum_ELBO(
          max_plate_nesting=0
      ).differentiable_loss(model, guide, data)
      actual_loss = JitTraceEnum_ELBO(
          max_plate_nesting=0
      ).differentiable_loss(model, guide, data)
      assert_equal(expected_loss, actual_loss)
  
      expected_grads = grad(
          expected_loss, [transition, means], allow_unused=True
      )
      actual_grads = grad(
          actual_loss, [transition, means], allow_unused=True
      )
      for e, a, name in zip(
          expected_grads, actual_grads, ["transition", "means"]
      ):
          assert_equal(e, a, msg="bad gradient for {}".format(name))
  
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 20 for calling a global function with example usage: test_traceenum_elbo(length) and returns: tuple