code: def test_mask_type(mask):
      p = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
      p_masked = p.mask(mask)
      if isinstance(mask, bool):
          mask = torch.tensor(mask)
  
      x = p.sample()
      actual = p_masked.log_prob(x)
      expected = p.log_prob(x) * mask.float()
      assert_equal(actual, expected)
  
      actual = p_masked.score_parts(x)
      expected = p.score_parts(x)
      for a, e in zip(actual, expected):
          if isinstance(e, torch.Tensor):
              e = e * mask.float()
          assert_equal(a, e)
  
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 3 for calling a global function with example usage: test_mask_type(mask) without return types