code: def test_svi_enum(plate_dim, enumerate1, enumerate2):
      pyro.clear_param_store()
      num_particles = 10
      q = pyro.param(
          "q", constant(0.75), constraint=constraints.unit_interval
      )
      p = 0.2693204236205713  # for which kl(Bernoulli(q), Bernoulli(p)) = 0.5
  
      def model():
          pyro.sample("x", dist.Bernoulli(p))
          for i in pyro.plate("plate", plate_dim):
              pyro.sample("y_{}".format(i), dist.Bernoulli(p))
  
      def guide():
          q = pyro.param("q")
          pyro.sample(
              "x", dist.Bernoulli(q), infer={"enumerate": enumerate1}
          )
          for i in pyro.plate("plate", plate_dim):
              pyro.sample(
                  "y_{}".format(i),
                  dist.Bernoulli(q),
                  infer={"enumerate": enumerate2},
              )
  
      kl = (1 + plate_dim) * kl_divergence(
          dist.Bernoulli(q), dist.Bernoulli(p)
      )
      expected_loss = kl.item()
      expected_grad = grad(kl, [q.unconstrained()])[0]
  
      inner_particles = 2
      outer_particles = num_particles // inner_particles
      elbo = TraceEnum_ELBO(
          max_plate_nesting=0,
          strict_enumeration_warning=any([enumerate1, enumerate2]),
          num_particles=inner_particles,
          ignore_jit_warnings=True,
      )
      actual_loss = (
          sum(
              elbo.loss_and_grads(model, guide)
              for i in range(outer_particles)
          )
          / outer_particles
      )
      actual_grad = q.unconstrained().grad / outer_particles
  
      assert_equal(
          actual_loss,
          expected_loss,
          prec=0.3,
          msg="".join(
              [
                  "\nexpected loss = {}".format(expected_loss),
                  "\n  actual loss = {}".format(actual_loss),
              ]
          ),
      )
      assert_equal(
          actual_grad,
          expected_grad,
          prec=0.5,
          msg="".join(
              [
                  "\nexpected grad = {}".format(
                      expected_grad.detach().cpu().numpy()
                  ),
                  "\n  actual grad = {}".format(
                      actual_grad.detach().cpu().numpy()
                  ),
              ]
          ),
      )
  
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 16 for calling a global function with example usage: test_svi_enum(plate_dim, enumerate1, enumerate2) without return types