code: def _check_traces(tr_pyro, tr_funsor):
      assert tr_pyro.nodes.keys() == tr_funsor.nodes.keys()
      tr_pyro.compute_log_prob()
      tr_funsor.compute_log_prob()
      tr_pyro.pack_tensors()
  
      symbol_to_name = {
          node["infer"]["_enumerate_symbol"]: name
          for name, node in tr_pyro.nodes.items()
          if node["type"] == "sample"
          and not node["is_observed"]
          and node["infer"].get("enumerate") == "parallel"
      }
      symbol_to_name.update(
          {
              symbol: name
              for name, symbol in tr_pyro.plate_to_symbol.items()
          }
      )
  
      if _NAMED_TEST_STRENGTH >= 1:
          # coarser check: enumeration requirements satisfied
          check_traceenum_requirements(tr_pyro, Trace())
          check_traceenum_requirements(tr_funsor, Trace())
          try:
              # coarser check: number of elements and squeezed shapes
              for name, pyro_node in tr_pyro.nodes.items():
                  if pyro_node["type"] != "sample":
                      continue
                  funsor_node = tr_funsor.nodes[name]
                  assert (
                      pyro_node["packed"]["log_prob"].numel()
                      == funsor_node["log_prob"].numel()
                  )
                  assert (
                      pyro_node["packed"]["log_prob"].shape
                      == funsor_node["log_prob"].squeeze().shape
                  )
                  assert frozenset(
                      f
                      for f in pyro_node["cond_indep_stack"]
                      if f.vectorized
                  ) == frozenset(
                      f
                      for f in funsor_node["cond_indep_stack"]
                      if f.vectorized
                  )
          except AssertionError:
              for name, pyro_node in tr_pyro.nodes.items():
                  if pyro_node["type"] != "sample":
                      continue
                  funsor_node = tr_funsor.nodes[name]
                  pyro_packed_shape = pyro_node["packed"][
                      "log_prob"
                  ].shape
                  funsor_packed_shape = (
                      funsor_node["log_prob"].squeeze().shape
                  )
                  if pyro_packed_shape != funsor_packed_shape:
                      err_str = "==> (dep mismatch) {}".format(name)
                  else:
                      err_str = name
                  print(
                      err_str,
                      "Pyro: {} vs Funsor: {}".format(
                          pyro_packed_shape, funsor_packed_shape
                      ),
                  )
              raise
  
      if _NAMED_TEST_STRENGTH >= 2:
          try:
              # medium check: unordered packed shapes match
              for name, pyro_node in tr_pyro.nodes.items():
                  if pyro_node["type"] != "sample":
                      continue
                  funsor_node = tr_funsor.nodes[name]
                  pyro_names = frozenset(
                      symbol_to_name[d]
                      for d in pyro_node["packed"][
                          "log_prob"
                      ]._pyro_dims
                  )
                  funsor_names = frozenset(
                      funsor_node["funsor"]["log_prob"].inputs
                  )
                  assert pyro_names == frozenset(
                      name.replace("__PARTICLES", "")
                      for name in funsor_names
                  )
          except AssertionError:
              for name, pyro_node in tr_pyro.nodes.items():
                  if pyro_node["type"] != "sample":
                      continue
                  funsor_node = tr_funsor.nodes[name]
                  pyro_names = frozenset(
                      symbol_to_name[d]
                      for d in pyro_node["packed"][
                          "log_prob"
                      ]._pyro_dims
                  )
                  funsor_names = frozenset(
                      funsor_node["funsor"]["log_prob"].inputs
                  )
                  if pyro_names != funsor_names:
                      err_str = "==> (packed mismatch) {}".format(name)
                  else:
                      err_str = name
                  print(
                      err_str,
                      "Pyro: {} vs Funsor: {}".format(
                          sorted(tuple(pyro_names)),
                          sorted(tuple(funsor_names)),
                      ),
                  )
              raise
  
      if _NAMED_TEST_STRENGTH >= 3:
          try:
              # finer check: exact match with unpacked Pyro shapes
              for name, pyro_node in tr_pyro.nodes.items():
                  if pyro_node["type"] != "sample":
                      continue
                  funsor_node = tr_funsor.nodes[name]
                  assert (
                      pyro_node["log_prob"].shape
                      == funsor_node["log_prob"].shape
                  )
                  assert (
                      pyro_node["value"].shape
                      == funsor_node["value"].shape
                  )
          except AssertionError:
              for name, pyro_node in tr_pyro.nodes.items():
                  if pyro_node["type"] != "sample":
                      continue
                  funsor_node = tr_funsor.nodes[name]
                  pyro_shape = pyro_node["log_prob"].shape
                  funsor_shape = funsor_node["log_prob"].shape
                  if pyro_shape != funsor_shape:
                      err_str = "==> (unpacked mismatch) {}".format(
                          name
                      )
                  else:
                      err_str = name
                  print(
                      err_str,
                      "Pyro: {} vs Funsor: {}".format(
                          pyro_shape, funsor_shape
                      ),
                  )
              raise
  
imports: import contextlib
  import logging
  import os
  from collections import defaultdict
  from queue import LifoQueue
  import pytest
  import torch
  from pyro.infer.enum import iter_discrete_escape, iter_discrete_extend
  from pyro.ops.indexing import Vindex
  from pyro.poutine import Trace
  from pyro.poutine.util import prune_subsample_sites
  from pyro.util import check_traceenum_requirements
label_desc: Use image node_id 2 for calling a global function with example usage: _check_traces(tr_pyro, tr_funsor) without return types