code: def test_distribution_3_simple(temperature):
      #  +---------------+
      #  |  z2 ---> x2   |
      #  |             2 |
      #  +---------------+
      num_particles = 10000
      data = torch.tensor([-1.0, 1.0])
  
      @infer.config_enumerate
      def model(z2=None):
          p = pyro.param("p", torch.tensor([0.25, 0.75]))
          loc = pyro.param("loc", torch.tensor([-1.0, 1.0]))
          with pyro.plate("data", 2):
              z2 = pyro.sample("z2", dist.Categorical(p), obs=z2)
              pyro.sample("x2", dist.Normal(loc[z2], 1.0), obs=data)
  
      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 = {
          (z20, z21): handlers.trace(model).get_trace(
              z2=torch.tensor([z20, z21])
          )
          for z20 in [0, 1]
          for z21 in [0, 1]
      }
  
      # Check joint posterior over (z2[0], z2[1]).
      actual_probs = torch.empty(2, 2)
      expected_probs = torch.empty(2, 2)
      for (z20, z21), tr in conditioned_traces.items():
          expected_probs[z20, z21] = tr.log_prob_sum().exp()
          actual_probs[z20, z21] = (
              (
                  (sampled_trace.nodes["z2"]["value"][..., :1] == z20)
                  & (sampled_trace.nodes["z2"]["value"][..., 1:] == z21)
              )
              .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.reshape(-1),
          actual_probs.reshape(-1),
          prec=1e-2,
      )
  
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 4 for calling a global function with example usage: test_distribution_3_simple(temperature) without return types