code: def _brute_force_compute_downstream_costs(
      model_trace, guide_trace, non_reparam_nodes  #
  ):
      guide_nodes = [
          x
          for x in guide_trace.nodes
          if guide_trace.nodes[x]["type"] == "sample"
      ]
      downstream_costs, downstream_guide_cost_nodes = {}, {}
      stacks = get_plate_stacks(model_trace)
  
      for node in guide_nodes:
          downstream_costs[node] = MultiFrameTensor(
              (
                  stacks[node],
                  model_trace.nodes[node]["log_prob"]
                  - guide_trace.nodes[node]["log_prob"],
              )
          )
          downstream_guide_cost_nodes[node] = set([node])
  
          descendants = guide_trace.successors(node)
  
          for desc in descendants:
              desc_mft = MultiFrameTensor(
                  (
                      stacks[desc],
                      model_trace.nodes[desc]["log_prob"]
                      - guide_trace.nodes[desc]["log_prob"],
                  )
              )
              downstream_costs[node].add(*desc_mft.items())
              downstream_guide_cost_nodes[node].update([desc])
  
      for site in non_reparam_nodes:
          children_in_model = set()
          for node in downstream_guide_cost_nodes[site]:
              children_in_model.update(model_trace.successors(node))
          children_in_model.difference_update(
              downstream_guide_cost_nodes[site]
          )
          for child in children_in_model:
              assert model_trace.nodes[child]["type"] == "sample"
              child_mft = MultiFrameTensor(
                  (stacks[child], model_trace.nodes[child]["log_prob"])
              )
              downstream_costs[site].add(*child_mft.items())
              downstream_guide_cost_nodes[site].update([child])
  
      for k in non_reparam_nodes:
          downstream_costs[k] = downstream_costs[k].sum_to(
              guide_trace.nodes[k]["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 1 for calling a global function with example usage: _brute_force_compute_downstream_costs(model_trace, guide_trace, non_reparam_nodes) and returns: downstream_costs,  downstream_guide_cost_nodes