code: def test_distribution_masked(infer, temperature):
      #      +-------+
      #  z --|--> x  |
      #      +-------+
      num_particles = 10000
      data = torch.tensor([1.0, 2.0, 3.0])
      mask = torch.tensor([True, False, False])
  
      @config_enumerate
      def model(num_particles=1, z=None):
          p = pyro.param("p", torch.tensor(0.25))
          with pyro.plate("num_particles", num_particles, dim=-2):
              z = pyro.sample("z", dist.Bernoulli(p), obs=z)
              logger.info("z.shape = {}".format(z.shape))
              with pyro.plate("data", 3), poutine.mask(mask=mask):
                  pyro.sample("x", dist.Normal(z, 1.0), obs=data)
  
      first_available_dim = -3
      sampled_model = infer(model, first_available_dim, temperature)
      sampled_trace = poutine.trace(sampled_model).get_trace(
          num_particles
      )
      conditioned_traces = {
          z: poutine.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"].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)
  
imports: import logging
  import math
  import warnings
  import pytest
  import torch
  from torch.autograd import grad
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.infer import TraceEnum_ELBO
  from pyro.infer.discrete import TraceEnumSample_ELBO, infer_discrete
  from pyro.infer.enum import config_enumerate
  from tests.common import assert_equal
label_desc: Use image node_id 7 for calling a global function with example usage: test_distribution_masked(infer, temperature) without return types