code: def test_kl_divergence():
      mask = torch.tensor([[0, 1], [1, 1]]).bool()
      p = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
      q = Normal(torch.randn(2, 2), torch.randn(2, 2).exp())
      expected = kl_divergence(p.to_event(2), q.to_event(2))
      actual = kl_divergence(
          p.mask(mask).to_event(2), q.mask(mask).to_event(2)
      ) + kl_divergence(
          p.mask(~mask).to_event(2), q.mask(~mask).to_event(2)
      )
      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 5 for calling a global function with example usage: test_kl_divergence() without return types