code: def _provenance_compute_downstream_costs(model_trace, guide_trace):
      # replicate the logic from TraceGraph_ELBO
      # additionally compute downstream_guide_cost_nodes
      downstream_costs = defaultdict(lambda: MultiFrameTensor())
      downstream_guide_cost_nodes = defaultdict(lambda: set())
      for name, site in model_trace.nodes.items():
          if site["type"] == "sample":
              # add the log_prob to each non-reparam sample site upstream
              for key in get_provenance(site["log_prob_sum"]):
                  downstream_costs[key].add(
                      (site["cond_indep_stack"], site["log_prob"])
                  )
                  downstream_guide_cost_nodes[key] |= {name}
      for name, site in guide_trace.nodes.items():
          if site["type"] == "sample":
              # add the -log_prob to each non-reparam sample site upstream
              for key in get_provenance(site["log_prob_sum"]):
                  downstream_costs[key].add(
                      (site["cond_indep_stack"], -site["log_prob"])
                  )
                  downstream_guide_cost_nodes[key] |= {name}
  
      for node, downstream_cost in downstream_costs.items():
          guide_site = guide_trace.nodes[node]
          downstream_costs[node] = downstream_cost.sum_to(
              guide_site["cond_indep_stack"]
          )
  
      return downstream_costs, downstream_guide_cost_nodes
  
imports: import math
  from collections import defaultdict
  import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.infer.tracegraph_elbo import TrackNonReparam, _compute_downstream_costs
  from pyro.infer.util import MultiFrameTensor, get_plate_stacks
  from pyro.ops.provenance import get_provenance
  from pyro.poutine.util import prune_subsample_sites
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: _provenance_compute_downstream_costs(model_trace, guide_trace) and returns: downstream_costs,  downstream_guide_cost_nodes