code: def test_kl_divergence_type(p_mask, q_mask):
      p = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
      q = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
      mask = (
          (torch.tensor(p_mask) if isinstance(p_mask, bool) else p_mask)
          & (
              torch.tensor(q_mask)
              if isinstance(q_mask, bool)
              else q_mask
          )
      ).expand(2, 2)
  
      expected = kl_divergence(p, q)
      expected[~mask] = 0
  
      actual = kl_divergence(p.mask(p_mask), q.mask(q_mask))
      if p_mask is False or q_mask is False:
          assert isinstance(actual, float) and actual == 0.0
      else:
          assert_equal(actual, expected)
  
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 6 for calling a global function with example usage: test_kl_divergence_type(p_mask, q_mask) without return types