code: def test_distribution_2(temperature):
      #       +--------+
      #  z1 --|--> x1  |
      #   |   |        |
      #   V   |        |
      #  z2 --|--> x2  |
      #       +--------+
      num_particles = 10000
      data = torch.tensor([[-1.0, -1.0, 0.0], [-1.0, 1.0, 1.0]])
  
      @infer.config_enumerate
      def model(z1=None, z2=None):
          p = pyro.param("p", torch.tensor([[0.25, 0.75], [0.1, 0.9]]))
          loc = pyro.param("loc", torch.tensor([-1.0, 1.0]))
          z1 = pyro.sample("z1", dist.Categorical(p[0]), obs=z1)
          z2 = pyro.sample("z2", dist.Categorical(p[z1]), obs=z2)
          logger.info("z1.shape = {}".format(z1.shape))
          logger.info("z2.shape = {}".format(z2.shape))
          with pyro.plate("data", 3):
              pyro.sample("x1", dist.Normal(loc[z1], 1.0), obs=data[0])
              pyro.sample("x2", dist.Normal(loc[z2], 1.0), obs=data[1])
  
      first_available_dim = -3
      vectorized_model = (
          model
          if temperature == 0
          else pyro.plate("particles", size=num_particles, dim=-2)(
              model
          )
      )
      sampled_model = infer.infer_discrete(
          vectorized_model, first_available_dim, temperature
      )
      sampled_trace = handlers.trace(sampled_model).get_trace()
      conditioned_traces = {
          (z1, z2): handlers.trace(model).get_trace(
              z1=torch.tensor(z1), z2=torch.tensor(z2)
          )
          for z1 in [0, 1]
          for z2 in [0, 1]
      }
  
      # Check joint posterior over (z1, z2).
      actual_probs = torch.empty(2, 2)
      expected_probs = torch.empty(2, 2)
      for (z1, z2), tr in conditioned_traces.items():
          expected_probs[z1, z2] = tr.log_prob_sum().exp()
          actual_probs[z1, z2] = (
              (
                  (sampled_trace.nodes["z1"]["value"] == z1)
                  & (sampled_trace.nodes["z2"]["value"] == z2)
              )
              .float()
              .mean()
          )
  
      if temperature:
          expected_probs = expected_probs / expected_probs.sum()
      else:
          expected_max, argmax = expected_probs.reshape(-1).max(0)
          actual_max = sampled_trace.log_prob_sum()
          assert_equal(expected_max.log(), actual_max, prec=1e-5)
          expected_probs[:] = 0
          expected_probs.reshape(-1)[argmax] = 1
      assert_equal(
          expected_probs,
          actual_probs,
          prec=1e-2 if temperature else 1e-5,
      )
  
file_path: pyro/tests/contrib/funsor/test_infer_discrete.py
imports: import logging
  import os
  import pyroapi
  import pytest
  import torch
  from pyro.infer.autoguide import AutoNormal
  from tests.common import assert_equal
label_desc: Use image node_id 3 for calling a global function with example usage: test_distribution_2(temperature) without return types