code: def test_compute_downstream_costs_plate_reuse(dim1, dim2):
      with TrackNonReparam():
          guide_trace = poutine.trace(
              plate_reuse_model_guide, graph_type="dense"
          ).get_trace(include_obs=False, dim1=dim1, dim2=dim2)
          model_trace = poutine.trace(
              poutine.replay(
                  plate_reuse_model_guide, trace=guide_trace
              ),
              graph_type="dense",
          ).get_trace(include_obs=True, dim1=dim1, dim2=dim2)
  
      guide_trace = prune_subsample_sites(guide_trace)
      model_trace = prune_subsample_sites(model_trace)
      model_trace.compute_log_prob()
      guide_trace.compute_log_prob()
  
      non_reparam_nodes = set(guide_trace.nonreparam_stochastic_nodes)
      dc, dc_nodes = _compute_downstream_costs(
          model_trace, guide_trace, non_reparam_nodes
      )
      dc_brute, dc_nodes_brute = _brute_force_compute_downstream_costs(
          model_trace, guide_trace, non_reparam_nodes
      )
      (
          dc_provenance,
          dc_nodes_provenance,
      ) = _provenance_compute_downstream_costs(model_trace, guide_trace)
  
      assert dc_nodes == dc_nodes_brute
  
      for name, nodes in dc_nodes_provenance.items():
          assert nodes.issubset(dc_nodes[name])
          if name == "c2":
              assert nodes == {"c2", "obs"}
          else:
              assert nodes == {name}
  
      for k in dc:
          assert guide_trace.nodes[k]["log_prob"].size() == dc[k].size()
          assert_equal(dc[k], dc_brute[k])
          # expected downstream cost provenance
          expected_dc_provenance = MultiFrameTensor(
              (
                  model_trace.nodes[k]["cond_indep_stack"],
                  model_trace.nodes[k]["log_prob"],
              ),
              (
                  guide_trace.nodes[k]["cond_indep_stack"],
                  -guide_trace.nodes[k]["log_prob"],
              ),
          )
          if k == "c2":
              expected_dc_provenance.add(
                  (
                      model_trace.nodes["obs"]["cond_indep_stack"],
                      model_trace.nodes["obs"]["log_prob"],
                  )
              )
          expected_dc_provenance = expected_dc_provenance.sum_to(
              guide_trace.nodes[k]["cond_indep_stack"]
          )
          assert_equal(dc_provenance[k], expected_dc_provenance)
  
      expected_c1 = (
          model_trace.nodes["c1"]["log_prob"]
          - guide_trace.nodes["c1"]["log_prob"]
      )
      expected_c1 += (
          model_trace.nodes["b1"]["log_prob"]
          - guide_trace.nodes["b1"]["log_prob"]
      ).sum()
      expected_c1 += (
          model_trace.nodes["c2"]["log_prob"]
          - guide_trace.nodes["c2"]["log_prob"]
      )
      expected_c1 += model_trace.nodes["obs"]["log_prob"]
      assert_equal(expected_c1, dc["c1"])
  
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 13 for calling a global function with example usage: test_compute_downstream_costs_plate_reuse(dim1, dim2) without return types