code: def test_mask(batch_dim, event_dim, mask_dim):
      # Construct base distribution.
      shape = torch.Size([2, 3, 4, 5, 6][: batch_dim + event_dim])
      batch_shape = shape[:batch_dim]
      mask_shape = batch_shape[batch_dim - mask_dim :]
      base_dist = Bernoulli(0.1).expand_by(shape).to_event(event_dim)
  
      # Construct masked distribution.
      mask = checker_mask(mask_shape)
      dist = base_dist.mask(mask)
  
      # Check shape.
      sample = base_dist.sample()
      assert dist.batch_shape == base_dist.batch_shape
      assert dist.event_shape == base_dist.event_shape
      assert sample.shape == sample.shape
      assert (
          dist.log_prob(sample).shape
          == base_dist.log_prob(sample).shape
      )
  
      # Check values.
      assert_equal(dist.mean, base_dist.mean)
      assert_equal(dist.variance, base_dist.variance)
      assert_equal(
          dist.log_prob(sample),
          scale_and_mask(base_dist.log_prob(sample), mask=mask),
      )
      assert_equal(
          dist.score_parts(sample),
          base_dist.score_parts(sample).scale_and_mask(mask=mask),
          prec=0,
      )
      if not dist.event_shape:
          assert_equal(
              dist.enumerate_support(), base_dist.enumerate_support()
          )
          assert_equal(
              dist.enumerate_support(expand=True),
              base_dist.enumerate_support(expand=True),
          )
          assert_equal(
              dist.enumerate_support(expand=False),
              base_dist.enumerate_support(expand=False),
          )
  
file_path: pyro/tests/distributions/test_mask.py
imports: import pytest
  import torch
  from torch import tensor
  from torch.distributions import kl_divergence
  from pyro.distributions.torch import Bernoulli, Normal
  from pyro.distributions.util import broadcast_shape, scale_and_mask
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_mask(batch_dim, event_dim, mask_dim) without return types