code: def test_compute_downstream_costs_plate_in_iplate(dim1):
      with TrackNonReparam():
          guide_trace = poutine.trace(
              nested_model_guide, graph_type="dense"
          ).get_trace(include_obs=False, dim1=dim1)
          model_trace = poutine.trace(
              poutine.replay(nested_model_guide, trace=guide_trace),
              graph_type="dense",
          ).get_trace(include_obs=True, dim1=dim1)
  
      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.startswith("c"):
              i = int(name[1:])
              assert nodes == {name, f"obs{i}"}
          else:
              assert nodes == {name}
  
      expected_c1 = (
          model_trace.nodes["c1"]["log_prob"]
          - guide_trace.nodes["c1"]["log_prob"]
      )
      expected_c1 += model_trace.nodes["obs1"]["log_prob"]
  
      expected_b1 = (
          model_trace.nodes["b1"]["log_prob"]
          - guide_trace.nodes["b1"]["log_prob"]
      )
      expected_b1 += (
          model_trace.nodes["c1"]["log_prob"]
          - guide_trace.nodes["c1"]["log_prob"]
      ).sum()
      expected_b1 += model_trace.nodes["obs1"]["log_prob"].sum()
  
      expected_c0 = (
          model_trace.nodes["c0"]["log_prob"]
          - guide_trace.nodes["c0"]["log_prob"]
      )
      expected_c0 += model_trace.nodes["obs0"]["log_prob"]
  
      expected_b0 = (
          model_trace.nodes["b0"]["log_prob"]
          - guide_trace.nodes["b0"]["log_prob"]
      )
      expected_b0 += (
          model_trace.nodes["c0"]["log_prob"]
          - guide_trace.nodes["c0"]["log_prob"]
      ).sum()
      expected_b0 += model_trace.nodes["obs0"]["log_prob"].sum()
  
      assert_equal(expected_c1, dc["c1"], prec=1.0e-6)
      assert_equal(expected_b1, dc["b1"], prec=1.0e-6)
      assert_equal(expected_c0, dc["c0"], prec=1.0e-6)
      assert_equal(expected_b0, dc["b0"], prec=1.0e-6)
  
      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.startswith("c"):
              i = int(k[1:])
              expected_dc_provenance.add(
                  (
                      model_trace.nodes[f"obs{i}"]["cond_indep_stack"],
                      model_trace.nodes[f"obs{i}"]["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)
  
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 9 for calling a global function with example usage: test_compute_downstream_costs_plate_in_iplate(dim1) without return types