code: def test_elbo_enumerate_plate_2(num_samples, num_masked, scale):
      #      +-----------------+
      #  x ----> y ----> z     |
      #      |               N |
      #      +-----------------+
      pyro.param(
          "guide_probs_x",
          torch.tensor([0.1, 0.9]),
          constraint=constraints.simplex,
      )
      pyro.param(
          "model_probs_x",
          torch.tensor([0.4, 0.6]),
          constraint=constraints.simplex,
      )
      pyro.param(
          "model_probs_y",
          torch.tensor([[0.75, 0.25], [0.55, 0.45]]),
          constraint=constraints.simplex,
      )
      pyro.param(
          "model_probs_z",
          torch.tensor([[0.3, 0.7], [0.2, 0.8]]),
          constraint=constraints.simplex,
      )
  
      def auto_model(data):
          probs_x = pyro.param("model_probs_x")
          probs_y = pyro.param("model_probs_y")
          probs_z = pyro.param("model_probs_z")
          x = pyro.sample("x", dist.Categorical(probs_x))
          with handlers.scale(scale=scale):
              with pyro.plate("data", len(data)):
                  if num_masked == num_samples:
                      y = pyro.sample(
                          "y",
                          dist.Categorical(probs_y[x]),
                          infer={"enumerate": "parallel"},
                      )
                      pyro.sample(
                          "z", dist.Categorical(probs_z[y]), obs=data
                      )
                  else:
                      with handlers.mask(
                          mask=torch.arange(num_samples) < num_masked
                      ):
                          y = pyro.sample(
                              "y",
                              dist.Categorical(probs_y[x]),
                              infer={"enumerate": "parallel"},
                          )
                          pyro.sample(
                              "z",
                              dist.Categorical(probs_z[y]),
                              obs=data,
                          )
  
      def hand_model(data):
          probs_x = pyro.param("model_probs_x")
          probs_y = pyro.param("model_probs_y")
          probs_z = pyro.param("model_probs_z")
          x = pyro.sample("x", dist.Categorical(probs_x))
          with handlers.scale(scale=scale):
              for i in pyro.plate("data", num_masked):
                  y = pyro.sample(
                      "y_{}".format(i),
                      dist.Categorical(probs_y[x]),
                      infer={"enumerate": "parallel"},
                  )
                  pyro.sample(
                      "z_{}".format(i),
                      dist.Categorical(probs_z[y]),
                      obs=data[i],
                  )
  
      @infer.config_enumerate
      def guide(data):
          probs_x = pyro.param("guide_probs_x")
          pyro.sample("x", dist.Categorical(probs_x))
  
      data = dist.Categorical(torch.tensor([0.3, 0.7])).sample(
          (num_samples,)
      )
      elbo = infer.TraceEnum_ELBO(max_plate_nesting=1)
      auto_loss = elbo.differentiable_loss(auto_model, guide, data)
      hand_loss = elbo.differentiable_loss(hand_model, guide, data)
      _check_loss_and_grads(hand_loss, auto_loss)
  
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 7 for calling a global function with example usage: test_elbo_enumerate_plate_2(num_samples, num_masked, scale) without return types