code: def test_elbo_enumerate_plates_6(scale):
      #         +----------+
      #         |      M=2 |
      #     a ----> b      |
      #     |   |   |      |
      #  +--|-------|--+   |
      #  |  V   |   V  |   |
      #  |  c ----> d  |   |
      #  |      |      |   |
      #  | N=2  +------|---+
      #  +-------------+
      # This tests different ways of mixing two independence contexts,
      # where each can be either sequential or vectorized plate.
      pyro.param(
          "probs_a",
          torch.tensor([0.45, 0.55]),
          constraint=constraints.simplex,
      )
      pyro.param(
          "probs_b",
          torch.tensor([[0.6, 0.4], [0.4, 0.6]]),
          constraint=constraints.simplex,
      )
      pyro.param(
          "probs_c",
          torch.tensor([[0.75, 0.25], [0.55, 0.45]]),
          constraint=constraints.simplex,
      )
      pyro.param(
          "probs_d",
          torch.tensor(
              [[[0.4, 0.6], [0.3, 0.7]], [[0.3, 0.7], [0.2, 0.8]]]
          ),
          constraint=constraints.simplex,
      )
  
      @infer.config_enumerate
      @handlers.scale(scale=scale)
      @handlers.trace
      def model_iplate_iplate(data):
          probs_a = pyro.param("probs_a")
          probs_b = pyro.param("probs_b")
          probs_c = pyro.param("probs_c")
          probs_d = pyro.param("probs_d")
          b_axis = pyro.plate("b_axis", 2)
          c_axis = pyro.plate("c_axis", 2)
          a = pyro.sample("a", dist.Categorical(probs_a))
          b = [
              pyro.sample(
                  "b_{}".format(i), dist.Categorical(probs_b[a])
              )
              for i in b_axis
          ]
          c = [
              pyro.sample(
                  "c_{}".format(j), dist.Categorical(probs_c[a])
              )
              for j in c_axis
          ]
          for i in b_axis:
              for j in c_axis:
                  pyro.sample(
                      "d_{}_{}".format(i, j),
                      dist.Categorical(Vindex(probs_d)[b[i], c[j]]),
                      obs=data[i, j],
                  )
  
      @infer.config_enumerate
      @handlers.scale(scale=scale)
      def model_iplate_plate(data):
          probs_a = pyro.param("probs_a")
          probs_b = pyro.param("probs_b")
          probs_c = pyro.param("probs_c")
          probs_d = pyro.param("probs_d")
          b_axis = pyro.plate("b_axis", 2)
          c_axis = pyro.plate("c_axis", 2)
          a = pyro.sample("a", dist.Categorical(probs_a))
          with c_axis:
              c = pyro.sample("c", dist.Categorical(probs_c[a]))
          for i in b_axis:
              b_i = pyro.sample(
                  "b_{}".format(i), dist.Categorical(probs_b[a])
              )
              with c_axis:
                  pyro.sample(
                      "d_{}".format(i),
                      dist.Categorical(Vindex(probs_d)[b_i, c]),
                      obs=data[i],
                  )
  
      @infer.config_enumerate
      @handlers.scale(scale=scale)
      @handlers.trace
      def model_plate_iplate(data):
          probs_a = pyro.param("probs_a")
          probs_b = pyro.param("probs_b")
          probs_c = pyro.param("probs_c")
          probs_d = pyro.param("probs_d")
          b_axis = pyro.plate("b_axis", 2)
          c_axis = pyro.plate("c_axis", 2)
          a = pyro.sample("a", dist.Categorical(probs_a))
          with b_axis:
              b = pyro.sample("b", dist.Categorical(probs_b[a]))
          c = [
              pyro.sample(
                  "c_{}".format(j), dist.Categorical(probs_c[a])
              )
              for j in c_axis
          ]
          with b_axis:
              for j in c_axis:
                  pyro.sample(
                      "d_{}".format(j),
                      dist.Categorical(Vindex(probs_d)[b, c[j]]),
                      obs=data[:, j],
                  )
  
      @infer.config_enumerate
      @handlers.scale(scale=scale)
      def model_plate_plate(data):
          probs_a = pyro.param("probs_a")
          probs_b = pyro.param("probs_b")
          probs_c = pyro.param("probs_c")
          probs_d = pyro.param("probs_d")
          b_axis = pyro.plate("b_axis", 2, dim=-1)
          c_axis = pyro.plate("c_axis", 2, dim=-2)
          a = pyro.sample("a", dist.Categorical(probs_a))
          with b_axis:
              b = pyro.sample("b", dist.Categorical(probs_b[a]))
          with c_axis:
              c = pyro.sample("c", dist.Categorical(probs_c[a]))
          with b_axis, c_axis:
              pyro.sample(
                  "d", dist.Categorical(Vindex(probs_d)[b, c]), obs=data
              )
  
      def guide(data):
          pass
  
      # Check that either one of the sequential plates can be promoted to be vectorized.
      data = torch.tensor([[0, 1], [0, 0]])
      elbo = infer.TraceEnum_ELBO(max_plate_nesting=0)
      loss_iplate_iplate = elbo.differentiable_loss(
          model_iplate_iplate, guide, data
      )
      elbo = infer.TraceEnum_ELBO(max_plate_nesting=1)
      loss_plate_iplate = elbo.differentiable_loss(
          model_plate_iplate, guide, data
      )
      loss_iplate_plate = elbo.differentiable_loss(
          model_iplate_plate, guide, data
      )
      _check_loss_and_grads(loss_iplate_iplate, loss_iplate_plate)
      _check_loss_and_grads(loss_iplate_iplate, loss_plate_iplate)
  
      # But promoting both to plates should result in an error.
      elbo = infer.TraceEnum_ELBO(max_plate_nesting=2)
      with pytest.raises(ValueError, match="intractable!"):
          elbo.differentiable_loss(model_plate_plate, guide, data)
  
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 18 for calling a global function with example usage: test_elbo_enumerate_plates_6(scale) without return types