code: def test_elbo_plate_plate(outer_dim, inner_dim):
      pyro.get_param_store().clear()
      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 = torch.distributions.kl.kl_divergence(
          torch.distributions.Categorical(q),
          torch.distributions.Categorical(p),
      )
      kl = (1 + outer_dim + inner_dim + outer_dim * inner_dim) * kl_node
      expected_loss = kl
      expected_grad = grad(kl, [q.unconstrained()])[0]
  
      elbo = infer.TraceEnum_ELBO(max_plate_nesting=2)
      actual_loss = elbo.differentiable_loss(model, guide)
      actual_grad = grad(actual_loss, [q.unconstrained()])[0]
  
      assert_equal(actual_loss, expected_loss, prec=1e-5)
      assert_equal(actual_grad, expected_grad, prec=1e-5)
  
imports: import logging
  import os
  import pyroapi
  import pytest
  import torch
  from torch.autograd import grad
  from torch.distributions import constraints
  from pyro.ops.indexing import Vindex
  from pyro.util import torch_isnan
  from tests.common import assert_equal, xfail_param
label_desc: Use image node_id 2 for calling a global function with example usage: test_elbo_plate_plate(outer_dim, inner_dim) without return types