code: def test_elbo_plate_plate(backend, outer_dim, inner_dim):
      with pyro_backend(backend):
          pyro.get_param_store().clear()
          num_particles = 1
          q = pyro.param(
              "q", torch.tensor([0.75, 0.25], requires_grad=True)
          )
          p = 0.2693204236205713  # for which kl(Categorical(q), Categorical(p)) = 0.5
          p = torch.tensor([p, 1 - p])
  
          def model():
              d = dist.Categorical(p)
              context1 = pyro.plate("outer", outer_dim, dim=-1)
              context2 = pyro.plate("inner", inner_dim, dim=-2)
              pyro.sample("w", d)
              with context1:
                  pyro.sample("x", d)
              with context2:
                  pyro.sample("y", d)
              with context1, context2:
                  pyro.sample("z", d)
  
          def guide():
              d = dist.Categorical(pyro.param("q"))
              context1 = pyro.plate("outer", outer_dim, dim=-1)
              context2 = pyro.plate("inner", inner_dim, dim=-2)
              pyro.sample("w", d, infer={"enumerate": "parallel"})
              with context1:
                  pyro.sample("x", d, infer={"enumerate": "parallel"})
              with context2:
                  pyro.sample("y", d, infer={"enumerate": "parallel"})
              with context1, context2:
                  pyro.sample("z", d, infer={"enumerate": "parallel"})
  
          kl_node = kl_divergence(
              torch.distributions.Categorical(funsor.to_data(q)),
              torch.distributions.Categorical(funsor.to_data(p)),
          )
          kl = (
              1 + outer_dim + inner_dim + outer_dim * inner_dim
          ) * kl_node
          expected_loss = kl
          expected_grad = grad(kl, [funsor.to_data(q)])[0]
  
          elbo = infer.TraceEnum_ELBO(
              num_particles=num_particles,
              vectorize_particles=True,
              strict_enumeration_warning=True,
          )
          elbo = elbo.differentiable_loss if backend == "pyro" else elbo
          actual_loss = funsor.to_data(elbo(model, guide))
          actual_loss.backward()
          actual_grad = funsor.to_data(pyro.param("q")).grad
  
          assert ops.allclose(actual_loss, expected_loss, atol=1e-5)
          assert ops.allclose(actual_grad, expected_grad, atol=1e-5)
  
imports: import warnings
  import pytest
  import funsor
  from funsor.testing import xfail_param
  from funsor.util import get_backend
label_desc: Use image node_id 18 for calling a global function with example usage: test_elbo_plate_plate(backend, outer_dim, inner_dim) without return types