code: def test_multi_frame_tensor():
      stacks = {}
      actual = MultiFrameTensor()
      tr = poutine.trace(xy_model).get_trace()
      for name, site in tr.nodes.items():
          if site["type"] == "sample":
              log_prob = site["fn"].log_prob(site["value"])
              stacks[name] = site["cond_indep_stack"]
              actual.add((site["cond_indep_stack"], log_prob))
  
      assert len(actual) == 4
  
      logp = math.log(0.5)
      expected = {
          "b": torch.ones(torch.Size()) * logp * (1 + 2 + 3 + 6),
          "bx": torch.ones(torch.Size((2,))) * logp * (1 + 1 + 3 + 3),
          "by": torch.ones(torch.Size((3, 1))) * logp * (1 + 2 + 1 + 2),
          "bxy": torch.ones(torch.Size((3, 2)))
          * logp
          * (1 + 1 + 1 + 1),
      }
      for name, expected_sum in expected.items():
          actual_sum = actual.sum_to(stacks[name])
          assert_equal(actual_sum, expected_sum, msg=name)
  
file_path: pyro/tests/infer/test_util.py
imports: import math
  import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.infer.importance import psis_diagnostic
  from pyro.infer.util import MultiFrameTensor
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_multi_frame_tensor() without return types