code: def test_distribution_masked(temperature):
      #      +-------+
      #  z --|--> x  |
      #      +-------+
      num_particles = 10000
      data = torch.tensor([1.0, 2.0, 3.0])
      mask = torch.tensor([True, False, False])
  
      @infer.config_enumerate
      def model(z=None):
          p = pyro.param("p", torch.tensor([0.75, 0.25]))
          z = pyro.sample("z", dist.Categorical(p), obs=z)
          logger.info("z.shape = {}".format(z.shape))
          with pyro.plate("data", 3), handlers.mask(mask=mask):
              pyro.sample(
                  "x", dist.Normal(z.type_as(data), 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 = {
          z: handlers.trace(model).get_trace(z=torch.tensor(z))
          for z in [0.0, 1.0]
      }
  
      # Check  posterior over z.
      actual_z_mean = (
          sampled_trace.nodes["z"]["value"].type_as(data).mean()
      )
      if temperature:
          expected_z_mean = 1 / (
              1
              + (
                  conditioned_traces[0].log_prob_sum()
                  - conditioned_traces[1].log_prob_sum()
              ).exp()
          )
      else:
          expected_z_mean = (
              conditioned_traces[1].log_prob_sum()
              > conditioned_traces[0].log_prob_sum()
          ).float()
      assert_equal(actual_z_mean, expected_z_mean, 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 10 for calling a global function with example usage: test_distribution_masked(temperature) without return types