We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Some pyro.contrib.funsor tests are failing in torch==1.11.0
This issue records tests of `pyro.contrib.funsor` that started failing in #3045. In that PR I've xfailed the following tests so as to unblock that PR.
- [ ] `python examples/contrib/funsor/hmm.py --num-steps=1 --truncate=10 --model=5 --funsor`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plate_1`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plate_4`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plate_6`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plate_7`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_2`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_3`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_4`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_5`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_6`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_7`
- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_8`

cc @eb8680 @ordabayevy 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pyro/contrib/funsor/infer/traceenum_elbo.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 import contextlib
5 
6 import funsor
7 from funsor.adjoint import AdjointTape
8 
9 from pyro.contrib.funsor import to_data, to_funsor
10 from pyro.contrib.funsor.handlers import enum, plate, replay, trace
11 from pyro.contrib.funsor.infer.elbo import ELBO, Jit_ELBO
12 from pyro.distributions.util import copy_docs_from
13 from pyro.infer import TraceEnum_ELBO as _OrigTraceEnum_ELBO
14 
15 
16 # Work around a bug in unfold_contraction_generic_tuple interacting with
17 # Approximate introduced in https://github.com/pyro-ppl/funsor/pull/488 .
18 # Once fixed, this can be replaced by funsor.optimizer.apply_optimizer().
19 def apply_optimizer(x):
20     with funsor.interpretations.normalize:
21         expr = funsor.interpreter.reinterpret(x)
22 
23     with funsor.optimizer.optimize_base:
24         return funsor.interpreter.reinterpret(expr)
25 
26 
27 def terms_from_trace(tr):
28     """Helper function to extract elbo components from execution traces."""
29     # data structure containing densities, measures, scales, and identification
30     # of free variables as either product (plate) variables or sum (measure) variables
31     terms = {
32         "log_factors": [],
33         "log_measures": [],
34         "scale": to_funsor(1.0),
35         "plate_vars": frozenset(),
36         "measure_vars": frozenset(),
37         "plate_to_step": dict(),
38     }
39     for name, node in tr.nodes.items():
40         # add markov dimensions to the plate_to_step dictionary
41         if node["type"] == "markov_chain":
42             terms["plate_to_step"][node["name"]] = node["value"]
43             # ensure previous step variables are added to measure_vars
44             for step in node["value"]:
45                 terms["measure_vars"] |= frozenset(
46                     {
47                         var
48                         for var in step[1:-1]
49                         if tr.nodes[var]["funsor"].get("log_measure", None) is not None
50                     }
51                 )
52         if (
53             node["type"] != "sample"
54             or type(node["fn"]).__name__ == "_Subsample"
55             or node["infer"].get("_do_not_score", False)
56         ):
57             continue
58         # grab plate dimensions from the cond_indep_stack
59         terms["plate_vars"] |= frozenset(
60             f.name for f in node["cond_indep_stack"] if f.vectorized
61         )
62         # grab the log-measure, found only at sites that are not replayed or observed
63         if node["funsor"].get("log_measure", None) is not None:
64             terms["log_measures"].append(node["funsor"]["log_measure"])
65             # sum (measure) variables: the fresh non-plate variables at a site
66             terms["measure_vars"] |= (
67                 frozenset(node["funsor"]["value"].inputs) | {name}
68             ) - terms["plate_vars"]
69         # grab the scale, assuming a common subsampling scale
70         if (
71             node.get("replay_active", False)
72             and set(node["funsor"]["log_prob"].inputs) & terms["measure_vars"]
73             and float(to_data(node["funsor"]["scale"])) != 1.0
74         ):
75             # model site that depends on enumerated variable: common scale
76             terms["scale"] = node["funsor"]["scale"]
77         else:  # otherwise: default scale behavior
78             node["funsor"]["log_prob"] = (
79                 node["funsor"]["log_prob"] * node["funsor"]["scale"]
80             )
81         # grab the log-density, found at all sites except those that are not replayed
82         if node["is_observed"] or not node.get("replay_skipped", False):
83             terms["log_factors"].append(node["funsor"]["log_prob"])
84     # add plate dimensions to the plate_to_step dictionary
85     terms["plate_to_step"].update(
86         {plate: terms["plate_to_step"].get(plate, {}) for plate in terms["plate_vars"]}
87     )
88     return terms
89 
90 
91 @copy_docs_from(_OrigTraceEnum_ELBO)
92 class TraceMarkovEnum_ELBO(ELBO):
93     def differentiable_loss(self, model, guide, *args, **kwargs):
94 
95         # get batched, enumerated, to_funsor-ed traces from the guide and model
96         with plate(
97             size=self.num_particles
98         ) if self.num_particles > 1 else contextlib.ExitStack(), enum(
99             first_available_dim=(-self.max_plate_nesting - 1)
100             if self.max_plate_nesting
101             else None
102         ):
103             guide_tr = trace(guide).get_trace(*args, **kwargs)
104             model_tr = trace(replay(model, trace=guide_tr)).get_trace(*args, **kwargs)
105 
106         # extract from traces all metadata that we will need to compute the elbo
107         guide_terms = terms_from_trace(guide_tr)
108         model_terms = terms_from_trace(model_tr)
109 
110         # guide side enumeration is not supported
111         if any(guide_terms["plate_to_step"].values()):
112             raise NotImplementedError(
113                 "TraceMarkovEnum_ELBO does not yet support guide side Markov enumeration"
114             )
115 
116         # build up a lazy expression for the elbo
117         with funsor.terms.lazy:
118             # identify and contract out auxiliary variables in the model with partial_sum_product
119             contracted_factors, uncontracted_factors = [], []
120             for f in model_terms["log_factors"]:
121                 if model_terms["measure_vars"].intersection(f.inputs):
122                     contracted_factors.append(f)
123                 else:
124                     uncontracted_factors.append(f)
125             # incorporate the effects of subsampling and handlers.scale through a common scale factor
126             markov_dims = frozenset(
127                 {plate for plate, step in model_terms["plate_to_step"].items() if step}
128             )
129             contracted_costs = [
130                 model_terms["scale"] * f
131                 for f in funsor.sum_product.dynamic_partial_sum_product(
132                     funsor.ops.logaddexp,
133                     funsor.ops.add,
134                     model_terms["log_measures"] + contracted_factors,
135                     plate_to_step=model_terms["plate_to_step"],
136                     eliminate=model_terms["measure_vars"] | markov_dims,
137                 )
138             ]
139 
140             costs = contracted_costs + uncontracted_factors  # model costs: logp
141             costs += [-f for f in guide_terms["log_factors"]]  # guide costs: -logq
142 
143             # finally, integrate out guide variables in the elbo and all plates
144             plate_vars = guide_terms["plate_vars"] | model_terms["plate_vars"]
145             elbo = to_funsor(0, output=funsor.Real)
146             for cost in costs:
147                 # compute the marginal logq in the guide corresponding to this cost term
148                 log_prob = funsor.sum_product.sum_product(
149                     funsor.ops.logaddexp,
150                     funsor.ops.add,
151                     guide_terms["log_measures"],
152                     plates=plate_vars,
153                     eliminate=(plate_vars | guide_terms["measure_vars"])
154                     - frozenset(cost.inputs),
155                 )
156                 # compute the expected cost term E_q[logp] or E_q[-logq] using the marginal logq for q
157                 elbo_term = funsor.Integrate(
158                     log_prob, cost, guide_terms["measure_vars"] & frozenset(cost.inputs)
159                 )
160                 elbo += elbo_term.reduce(
161                     funsor.ops.add, plate_vars & frozenset(cost.inputs)
162                 )
163 
164         # evaluate the elbo, using memoize to share tensor computation where possible
165         with funsor.interpretations.memoize():
166             return -to_data(apply_optimizer(elbo))
167 
168 
169 @copy_docs_from(_OrigTraceEnum_ELBO)
170 class TraceEnum_ELBO(ELBO):
171     def differentiable_loss(self, model, guide, *args, **kwargs):
172 
173         # get batched, enumerated, to_funsor-ed traces from the guide and model
174         with plate(
175             size=self.num_particles
176         ) if self.num_particles > 1 else contextlib.ExitStack(), enum(
177             first_available_dim=(-self.max_plate_nesting - 1)
178             if self.max_plate_nesting
179             else None
180         ):
181             guide_tr = trace(guide).get_trace(*args, **kwargs)
182             model_tr = trace(replay(model, trace=guide_tr)).get_trace(*args, **kwargs)
183 
184         # extract from traces all metadata that we will need to compute the elbo
185         guide_terms = terms_from_trace(guide_tr)
186         model_terms = terms_from_trace(model_tr)
187 
188         # build up a lazy expression for the elbo
189         with funsor.terms.lazy:
190             # identify and contract out auxiliary variables in the model with partial_sum_product
191             contracted_factors, uncontracted_factors = [], []
192             for f in model_terms["log_factors"]:
193                 if model_terms["measure_vars"].intersection(f.inputs):
194                     contracted_factors.append(f)
195                 else:
196                     uncontracted_factors.append(f)
197             # incorporate the effects of subsampling and handlers.scale through a common scale factor
198             contracted_costs = [
199                 model_terms["scale"] * f
200                 for f in funsor.sum_product.partial_sum_product(
201                     funsor.ops.logaddexp,
202                     funsor.ops.add,
203                     model_terms["log_measures"] + contracted_factors,
204                     plates=model_terms["plate_vars"],
205                     eliminate=model_terms["measure_vars"],
206                 )
207             ]
208 
209             # accumulate costs from model (logp) and guide (-logq)
210             costs = contracted_costs + uncontracted_factors  # model costs: logp
211             costs += [-f for f in guide_terms["log_factors"]]  # guide costs: -logq
212 
213             # compute expected cost
214             # Cf. pyro.infer.util.Dice.compute_expectation()
215             # https://github.com/pyro-ppl/pyro/blob/0.3.0/pyro/infer/util.py#L212
216             # TODO Replace this with funsor.Expectation
217             plate_vars = guide_terms["plate_vars"] | model_terms["plate_vars"]
218             # compute the marginal logq in the guide corresponding to each cost term
219             targets = dict()
220             for cost in costs:
221                 input_vars = frozenset(cost.inputs)
222                 if input_vars not in targets:
223                     targets[input_vars] = funsor.Tensor(
224                         funsor.ops.new_zeros(
225                             funsor.tensor.get_default_prototype(),
226                             tuple(v.size for v in cost.inputs.values()),
227                         ),
228                         cost.inputs,
229                         cost.dtype,
230                     )
231             with AdjointTape() as tape:
232                 logzq = funsor.sum_product.sum_product(
233                     funsor.ops.logaddexp,
234                     funsor.ops.add,
235                     guide_terms["log_measures"] + list(targets.values()),
236                     plates=plate_vars,
237                     eliminate=(plate_vars | guide_terms["measure_vars"]),
238                 )
239             marginals = tape.adjoint(
240                 funsor.ops.logaddexp, funsor.ops.add, logzq, tuple(targets.values())
241             )
242             # finally, integrate out guide variables in the elbo and all plates
243             elbo = to_funsor(0, output=funsor.Real)
244             for cost in costs:
245                 target = targets[frozenset(cost.inputs)]
246                 logzq_local = marginals[target].reduce(
247                     funsor.ops.logaddexp, frozenset(cost.inputs) - plate_vars
248                 )
249                 log_prob = marginals[target] - logzq_local
250                 elbo_term = funsor.Integrate(
251                     log_prob,
252                     cost,
253                     guide_terms["measure_vars"] & frozenset(log_prob.inputs),
254                 )
255                 elbo += elbo_term.reduce(
256                     funsor.ops.add, plate_vars & frozenset(cost.inputs)
257                 )
258 
259         # evaluate the elbo, using memoize to share tensor computation where possible
260         with funsor.interpretations.memoize():
261             return -to_data(apply_optimizer(elbo))
262 
263 
264 class JitTraceEnum_ELBO(Jit_ELBO, TraceEnum_ELBO):
265     pass
266 
267 
268 class JitTraceMarkovEnum_ELBO(Jit_ELBO, TraceMarkovEnum_ELBO):
269     pass
270
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pyro/contrib/funsor/infer/traceenum_elbo.py b/pyro/contrib/funsor/infer/traceenum_elbo.py
--- a/pyro/contrib/funsor/infer/traceenum_elbo.py
+++ b/pyro/contrib/funsor/infer/traceenum_elbo.py
@@ -5,6 +5,7 @@
 
 import funsor
 from funsor.adjoint import AdjointTape
+from funsor.sum_product import _partition
 
 from pyro.contrib.funsor import to_data, to_funsor
 from pyro.contrib.funsor.handlers import enum, plate, replay, trace
@@ -194,17 +195,28 @@
                     contracted_factors.append(f)
                 else:
                     uncontracted_factors.append(f)
+            contracted_costs = []
             # incorporate the effects of subsampling and handlers.scale through a common scale factor
-            contracted_costs = [
-                model_terms["scale"] * f
+            for group_factors, group_vars in _partition(
+                model_terms["log_measures"] + contracted_factors,
+                model_terms["measure_vars"],
+            ):
+                group_factor_vars = frozenset().union(
+                    *[f.inputs for f in group_factors]
+                )
+                group_plates = model_terms["plate_vars"] & group_factor_vars
+                outermost_plates = frozenset.intersection(
+                    *(frozenset(f.inputs) & group_plates for f in group_factors)
+                )
+                elim_plates = group_plates - outermost_plates
                 for f in funsor.sum_product.partial_sum_product(
                     funsor.ops.logaddexp,
                     funsor.ops.add,
-                    model_terms["log_measures"] + contracted_factors,
-                    plates=model_terms["plate_vars"],
-                    eliminate=model_terms["measure_vars"],
-                )
-            ]
+                    group_factors,
+                    plates=group_plates,
+                    eliminate=group_vars | elim_plates,
+                ):
+                    contracted_costs.append(model_terms["scale"] * f)
 
             # accumulate costs from model (logp) and guide (-logq)
             costs = contracted_costs + uncontracted_factors  # model costs: logp

{"golden_diff": "diff --git a/pyro/contrib/funsor/infer/traceenum_elbo.py b/pyro/contrib/funsor/infer/traceenum_elbo.py\n--- a/pyro/contrib/funsor/infer/traceenum_elbo.py\n+++ b/pyro/contrib/funsor/infer/traceenum_elbo.py\n@@ -5,6 +5,7 @@\n \n import funsor\n from funsor.adjoint import AdjointTape\n+from funsor.sum_product import _partition\n \n from pyro.contrib.funsor import to_data, to_funsor\n from pyro.contrib.funsor.handlers import enum, plate, replay, trace\n@@ -194,17 +195,28 @@\n                     contracted_factors.append(f)\n                 else:\n                     uncontracted_factors.append(f)\n+            contracted_costs = []\n             # incorporate the effects of subsampling and handlers.scale through a common scale factor\n-            contracted_costs = [\n-                model_terms[\"scale\"] * f\n+            for group_factors, group_vars in _partition(\n+                model_terms[\"log_measures\"] + contracted_factors,\n+                model_terms[\"measure_vars\"],\n+            ):\n+                group_factor_vars = frozenset().union(\n+                    *[f.inputs for f in group_factors]\n+                )\n+                group_plates = model_terms[\"plate_vars\"] & group_factor_vars\n+                outermost_plates = frozenset.intersection(\n+                    *(frozenset(f.inputs) & group_plates for f in group_factors)\n+                )\n+                elim_plates = group_plates - outermost_plates\n                 for f in funsor.sum_product.partial_sum_product(\n                     funsor.ops.logaddexp,\n                     funsor.ops.add,\n-                    model_terms[\"log_measures\"] + contracted_factors,\n-                    plates=model_terms[\"plate_vars\"],\n-                    eliminate=model_terms[\"measure_vars\"],\n-                )\n-            ]\n+                    group_factors,\n+                    plates=group_plates,\n+                    eliminate=group_vars | elim_plates,\n+                ):\n+                    contracted_costs.append(model_terms[\"scale\"] * f)\n \n             # accumulate costs from model (logp) and guide (-logq)\n             costs = contracted_costs + uncontracted_factors  # model costs: logp\n", "issue": "Some pyro.contrib.funsor tests are failing in torch==1.11.0\nThis issue records tests of `pyro.contrib.funsor` that started failing in #3045. In that PR I've xfailed the following tests so as to unblock that PR.\r\n- [ ] `python examples/contrib/funsor/hmm.py --num-steps=1 --truncate=10 --model=5 --funsor`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plate_1`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plate_4`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plate_6`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plate_7`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_2`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_3`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_4`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_5`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_6`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_7`\r\n- [ ] `pytest -x tests/contrib/funsor/test_enum_funsor.py::test_elbo_enumerate_plates_8`\r\n\r\ncc @eb8680 @ordabayevy \n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nimport contextlib\n\nimport funsor\nfrom funsor.adjoint import AdjointTape\n\nfrom pyro.contrib.funsor import to_data, to_funsor\nfrom pyro.contrib.funsor.handlers import enum, plate, replay, trace\nfrom pyro.contrib.funsor.infer.elbo import ELBO, Jit_ELBO\nfrom pyro.distributions.util import copy_docs_from\nfrom pyro.infer import TraceEnum_ELBO as _OrigTraceEnum_ELBO\n\n\n# Work around a bug in unfold_contraction_generic_tuple interacting with\n# Approximate introduced in https://github.com/pyro-ppl/funsor/pull/488 .\n# Once fixed, this can be replaced by funsor.optimizer.apply_optimizer().\ndef apply_optimizer(x):\n    with funsor.interpretations.normalize:\n        expr = funsor.interpreter.reinterpret(x)\n\n    with funsor.optimizer.optimize_base:\n        return funsor.interpreter.reinterpret(expr)\n\n\ndef terms_from_trace(tr):\n    \"\"\"Helper function to extract elbo components from execution traces.\"\"\"\n    # data structure containing densities, measures, scales, and identification\n    # of free variables as either product (plate) variables or sum (measure) variables\n    terms = {\n        \"log_factors\": [],\n        \"log_measures\": [],\n        \"scale\": to_funsor(1.0),\n        \"plate_vars\": frozenset(),\n        \"measure_vars\": frozenset(),\n        \"plate_to_step\": dict(),\n    }\n    for name, node in tr.nodes.items():\n        # add markov dimensions to the plate_to_step dictionary\n        if node[\"type\"] == \"markov_chain\":\n            terms[\"plate_to_step\"][node[\"name\"]] = node[\"value\"]\n            # ensure previous step variables are added to measure_vars\n            for step in node[\"value\"]:\n                terms[\"measure_vars\"] |= frozenset(\n                    {\n                        var\n                        for var in step[1:-1]\n                        if tr.nodes[var][\"funsor\"].get(\"log_measure\", None) is not None\n                    }\n                )\n        if (\n            node[\"type\"] != \"sample\"\n            or type(node[\"fn\"]).__name__ == \"_Subsample\"\n            or node[\"infer\"].get(\"_do_not_score\", False)\n        ):\n            continue\n        # grab plate dimensions from the cond_indep_stack\n        terms[\"plate_vars\"] |= frozenset(\n            f.name for f in node[\"cond_indep_stack\"] if f.vectorized\n        )\n        # grab the log-measure, found only at sites that are not replayed or observed\n        if node[\"funsor\"].get(\"log_measure\", None) is not None:\n            terms[\"log_measures\"].append(node[\"funsor\"][\"log_measure\"])\n            # sum (measure) variables: the fresh non-plate variables at a site\n            terms[\"measure_vars\"] |= (\n                frozenset(node[\"funsor\"][\"value\"].inputs) | {name}\n            ) - terms[\"plate_vars\"]\n        # grab the scale, assuming a common subsampling scale\n        if (\n            node.get(\"replay_active\", False)\n            and set(node[\"funsor\"][\"log_prob\"].inputs) & terms[\"measure_vars\"]\n            and float(to_data(node[\"funsor\"][\"scale\"])) != 1.0\n        ):\n            # model site that depends on enumerated variable: common scale\n            terms[\"scale\"] = node[\"funsor\"][\"scale\"]\n        else:  # otherwise: default scale behavior\n            node[\"funsor\"][\"log_prob\"] = (\n                node[\"funsor\"][\"log_prob\"] * node[\"funsor\"][\"scale\"]\n            )\n        # grab the log-density, found at all sites except those that are not replayed\n        if node[\"is_observed\"] or not node.get(\"replay_skipped\", False):\n            terms[\"log_factors\"].append(node[\"funsor\"][\"log_prob\"])\n    # add plate dimensions to the plate_to_step dictionary\n    terms[\"plate_to_step\"].update(\n        {plate: terms[\"plate_to_step\"].get(plate, {}) for plate in terms[\"plate_vars\"]}\n    )\n    return terms\n\n\n@copy_docs_from(_OrigTraceEnum_ELBO)\nclass TraceMarkovEnum_ELBO(ELBO):\n    def differentiable_loss(self, model, guide, *args, **kwargs):\n\n        # get batched, enumerated, to_funsor-ed traces from the guide and model\n        with plate(\n            size=self.num_particles\n        ) if self.num_particles > 1 else contextlib.ExitStack(), enum(\n            first_available_dim=(-self.max_plate_nesting - 1)\n            if self.max_plate_nesting\n            else None\n        ):\n            guide_tr = trace(guide).get_trace(*args, **kwargs)\n            model_tr = trace(replay(model, trace=guide_tr)).get_trace(*args, **kwargs)\n\n        # extract from traces all metadata that we will need to compute the elbo\n        guide_terms = terms_from_trace(guide_tr)\n        model_terms = terms_from_trace(model_tr)\n\n        # guide side enumeration is not supported\n        if any(guide_terms[\"plate_to_step\"].values()):\n            raise NotImplementedError(\n                \"TraceMarkovEnum_ELBO does not yet support guide side Markov enumeration\"\n            )\n\n        # build up a lazy expression for the elbo\n        with funsor.terms.lazy:\n            # identify and contract out auxiliary variables in the model with partial_sum_product\n            contracted_factors, uncontracted_factors = [], []\n            for f in model_terms[\"log_factors\"]:\n                if model_terms[\"measure_vars\"].intersection(f.inputs):\n                    contracted_factors.append(f)\n                else:\n                    uncontracted_factors.append(f)\n            # incorporate the effects of subsampling and handlers.scale through a common scale factor\n            markov_dims = frozenset(\n                {plate for plate, step in model_terms[\"plate_to_step\"].items() if step}\n            )\n            contracted_costs = [\n                model_terms[\"scale\"] * f\n                for f in funsor.sum_product.dynamic_partial_sum_product(\n                    funsor.ops.logaddexp,\n                    funsor.ops.add,\n                    model_terms[\"log_measures\"] + contracted_factors,\n                    plate_to_step=model_terms[\"plate_to_step\"],\n                    eliminate=model_terms[\"measure_vars\"] | markov_dims,\n                )\n            ]\n\n            costs = contracted_costs + uncontracted_factors  # model costs: logp\n            costs += [-f for f in guide_terms[\"log_factors\"]]  # guide costs: -logq\n\n            # finally, integrate out guide variables in the elbo and all plates\n            plate_vars = guide_terms[\"plate_vars\"] | model_terms[\"plate_vars\"]\n            elbo = to_funsor(0, output=funsor.Real)\n            for cost in costs:\n                # compute the marginal logq in the guide corresponding to this cost term\n                log_prob = funsor.sum_product.sum_product(\n                    funsor.ops.logaddexp,\n                    funsor.ops.add,\n                    guide_terms[\"log_measures\"],\n                    plates=plate_vars,\n                    eliminate=(plate_vars | guide_terms[\"measure_vars\"])\n                    - frozenset(cost.inputs),\n                )\n                # compute the expected cost term E_q[logp] or E_q[-logq] using the marginal logq for q\n                elbo_term = funsor.Integrate(\n                    log_prob, cost, guide_terms[\"measure_vars\"] & frozenset(cost.inputs)\n                )\n                elbo += elbo_term.reduce(\n                    funsor.ops.add, plate_vars & frozenset(cost.inputs)\n                )\n\n        # evaluate the elbo, using memoize to share tensor computation where possible\n        with funsor.interpretations.memoize():\n            return -to_data(apply_optimizer(elbo))\n\n\n@copy_docs_from(_OrigTraceEnum_ELBO)\nclass TraceEnum_ELBO(ELBO):\n    def differentiable_loss(self, model, guide, *args, **kwargs):\n\n        # get batched, enumerated, to_funsor-ed traces from the guide and model\n        with plate(\n            size=self.num_particles\n        ) if self.num_particles > 1 else contextlib.ExitStack(), enum(\n            first_available_dim=(-self.max_plate_nesting - 1)\n            if self.max_plate_nesting\n            else None\n        ):\n            guide_tr = trace(guide).get_trace(*args, **kwargs)\n            model_tr = trace(replay(model, trace=guide_tr)).get_trace(*args, **kwargs)\n\n        # extract from traces all metadata that we will need to compute the elbo\n        guide_terms = terms_from_trace(guide_tr)\n        model_terms = terms_from_trace(model_tr)\n\n        # build up a lazy expression for the elbo\n        with funsor.terms.lazy:\n            # identify and contract out auxiliary variables in the model with partial_sum_product\n            contracted_factors, uncontracted_factors = [], []\n            for f in model_terms[\"log_factors\"]:\n                if model_terms[\"measure_vars\"].intersection(f.inputs):\n                    contracted_factors.append(f)\n                else:\n                    uncontracted_factors.append(f)\n            # incorporate the effects of subsampling and handlers.scale through a common scale factor\n            contracted_costs = [\n                model_terms[\"scale\"] * f\n                for f in funsor.sum_product.partial_sum_product(\n                    funsor.ops.logaddexp,\n                    funsor.ops.add,\n                    model_terms[\"log_measures\"] + contracted_factors,\n                    plates=model_terms[\"plate_vars\"],\n                    eliminate=model_terms[\"measure_vars\"],\n                )\n            ]\n\n            # accumulate costs from model (logp) and guide (-logq)\n            costs = contracted_costs + uncontracted_factors  # model costs: logp\n            costs += [-f for f in guide_terms[\"log_factors\"]]  # guide costs: -logq\n\n            # compute expected cost\n            # Cf. pyro.infer.util.Dice.compute_expectation()\n            # https://github.com/pyro-ppl/pyro/blob/0.3.0/pyro/infer/util.py#L212\n            # TODO Replace this with funsor.Expectation\n            plate_vars = guide_terms[\"plate_vars\"] | model_terms[\"plate_vars\"]\n            # compute the marginal logq in the guide corresponding to each cost term\n            targets = dict()\n            for cost in costs:\n                input_vars = frozenset(cost.inputs)\n                if input_vars not in targets:\n                    targets[input_vars] = funsor.Tensor(\n                        funsor.ops.new_zeros(\n                            funsor.tensor.get_default_prototype(),\n                            tuple(v.size for v in cost.inputs.values()),\n                        ),\n                        cost.inputs,\n                        cost.dtype,\n                    )\n            with AdjointTape() as tape:\n                logzq = funsor.sum_product.sum_product(\n                    funsor.ops.logaddexp,\n                    funsor.ops.add,\n                    guide_terms[\"log_measures\"] + list(targets.values()),\n                    plates=plate_vars,\n                    eliminate=(plate_vars | guide_terms[\"measure_vars\"]),\n                )\n            marginals = tape.adjoint(\n                funsor.ops.logaddexp, funsor.ops.add, logzq, tuple(targets.values())\n            )\n            # finally, integrate out guide variables in the elbo and all plates\n            elbo = to_funsor(0, output=funsor.Real)\n            for cost in costs:\n                target = targets[frozenset(cost.inputs)]\n                logzq_local = marginals[target].reduce(\n                    funsor.ops.logaddexp, frozenset(cost.inputs) - plate_vars\n                )\n                log_prob = marginals[target] - logzq_local\n                elbo_term = funsor.Integrate(\n                    log_prob,\n                    cost,\n                    guide_terms[\"measure_vars\"] & frozenset(log_prob.inputs),\n                )\n                elbo += elbo_term.reduce(\n                    funsor.ops.add, plate_vars & frozenset(cost.inputs)\n                )\n\n        # evaluate the elbo, using memoize to share tensor computation where possible\n        with funsor.interpretations.memoize():\n            return -to_data(apply_optimizer(elbo))\n\n\nclass JitTraceEnum_ELBO(Jit_ELBO, TraceEnum_ELBO):\n    pass\n\n\nclass JitTraceMarkovEnum_ELBO(Jit_ELBO, TraceMarkovEnum_ELBO):\n    pass\n", "path": "pyro/contrib/funsor/infer/traceenum_elbo.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nimport contextlib\n\nimport funsor\nfrom funsor.adjoint import AdjointTape\nfrom funsor.sum_product import _partition\n\nfrom pyro.contrib.funsor import to_data, to_funsor\nfrom pyro.contrib.funsor.handlers import enum, plate, replay, trace\nfrom pyro.contrib.funsor.infer.elbo import ELBO, Jit_ELBO\nfrom pyro.distributions.util import copy_docs_from\nfrom pyro.infer import TraceEnum_ELBO as _OrigTraceEnum_ELBO\n\n\n# Work around a bug in unfold_contraction_generic_tuple interacting with\n# Approximate introduced in https://github.com/pyro-ppl/funsor/pull/488 .\n# Once fixed, this can be replaced by funsor.optimizer.apply_optimizer().\ndef apply_optimizer(x):\n    with funsor.interpretations.normalize:\n        expr = funsor.interpreter.reinterpret(x)\n\n    with funsor.optimizer.optimize_base:\n        return funsor.interpreter.reinterpret(expr)\n\n\ndef terms_from_trace(tr):\n    \"\"\"Helper function to extract elbo components from execution traces.\"\"\"\n    # data structure containing densities, measures, scales, and identification\n    # of free variables as either product (plate) variables or sum (measure) variables\n    terms = {\n        \"log_factors\": [],\n        \"log_measures\": [],\n        \"scale\": to_funsor(1.0),\n        \"plate_vars\": frozenset(),\n        \"measure_vars\": frozenset(),\n        \"plate_to_step\": dict(),\n    }\n    for name, node in tr.nodes.items():\n        # add markov dimensions to the plate_to_step dictionary\n        if node[\"type\"] == \"markov_chain\":\n            terms[\"plate_to_step\"][node[\"name\"]] = node[\"value\"]\n            # ensure previous step variables are added to measure_vars\n            for step in node[\"value\"]:\n                terms[\"measure_vars\"] |= frozenset(\n                    {\n                        var\n                        for var in step[1:-1]\n                        if tr.nodes[var][\"funsor\"].get(\"log_measure\", None) is not None\n                    }\n                )\n        if (\n            node[\"type\"] != \"sample\"\n            or type(node[\"fn\"]).__name__ == \"_Subsample\"\n            or node[\"infer\"].get(\"_do_not_score\", False)\n        ):\n            continue\n        # grab plate dimensions from the cond_indep_stack\n        terms[\"plate_vars\"] |= frozenset(\n            f.name for f in node[\"cond_indep_stack\"] if f.vectorized\n        )\n        # grab the log-measure, found only at sites that are not replayed or observed\n        if node[\"funsor\"].get(\"log_measure\", None) is not None:\n            terms[\"log_measures\"].append(node[\"funsor\"][\"log_measure\"])\n            # sum (measure) variables: the fresh non-plate variables at a site\n            terms[\"measure_vars\"] |= (\n                frozenset(node[\"funsor\"][\"value\"].inputs) | {name}\n            ) - terms[\"plate_vars\"]\n        # grab the scale, assuming a common subsampling scale\n        if (\n            node.get(\"replay_active\", False)\n            and set(node[\"funsor\"][\"log_prob\"].inputs) & terms[\"measure_vars\"]\n            and float(to_data(node[\"funsor\"][\"scale\"])) != 1.0\n        ):\n            # model site that depends on enumerated variable: common scale\n            terms[\"scale\"] = node[\"funsor\"][\"scale\"]\n        else:  # otherwise: default scale behavior\n            node[\"funsor\"][\"log_prob\"] = (\n                node[\"funsor\"][\"log_prob\"] * node[\"funsor\"][\"scale\"]\n            )\n        # grab the log-density, found at all sites except those that are not replayed\n        if node[\"is_observed\"] or not node.get(\"replay_skipped\", False):\n            terms[\"log_factors\"].append(node[\"funsor\"][\"log_prob\"])\n    # add plate dimensions to the plate_to_step dictionary\n    terms[\"plate_to_step\"].update(\n        {plate: terms[\"plate_to_step\"].get(plate, {}) for plate in terms[\"plate_vars\"]}\n    )\n    return terms\n\n\n@copy_docs_from(_OrigTraceEnum_ELBO)\nclass TraceMarkovEnum_ELBO(ELBO):\n    def differentiable_loss(self, model, guide, *args, **kwargs):\n\n        # get batched, enumerated, to_funsor-ed traces from the guide and model\n        with plate(\n            size=self.num_particles\n        ) if self.num_particles > 1 else contextlib.ExitStack(), enum(\n            first_available_dim=(-self.max_plate_nesting - 1)\n            if self.max_plate_nesting\n            else None\n        ):\n            guide_tr = trace(guide).get_trace(*args, **kwargs)\n            model_tr = trace(replay(model, trace=guide_tr)).get_trace(*args, **kwargs)\n\n        # extract from traces all metadata that we will need to compute the elbo\n        guide_terms = terms_from_trace(guide_tr)\n        model_terms = terms_from_trace(model_tr)\n\n        # guide side enumeration is not supported\n        if any(guide_terms[\"plate_to_step\"].values()):\n            raise NotImplementedError(\n                \"TraceMarkovEnum_ELBO does not yet support guide side Markov enumeration\"\n            )\n\n        # build up a lazy expression for the elbo\n        with funsor.terms.lazy:\n            # identify and contract out auxiliary variables in the model with partial_sum_product\n            contracted_factors, uncontracted_factors = [], []\n            for f in model_terms[\"log_factors\"]:\n                if model_terms[\"measure_vars\"].intersection(f.inputs):\n                    contracted_factors.append(f)\n                else:\n                    uncontracted_factors.append(f)\n            # incorporate the effects of subsampling and handlers.scale through a common scale factor\n            markov_dims = frozenset(\n                {plate for plate, step in model_terms[\"plate_to_step\"].items() if step}\n            )\n            contracted_costs = [\n                model_terms[\"scale\"] * f\n                for f in funsor.sum_product.dynamic_partial_sum_product(\n                    funsor.ops.logaddexp,\n                    funsor.ops.add,\n                    model_terms[\"log_measures\"] + contracted_factors,\n                    plate_to_step=model_terms[\"plate_to_step\"],\n                    eliminate=model_terms[\"measure_vars\"] | markov_dims,\n                )\n            ]\n\n            costs = contracted_costs + uncontracted_factors  # model costs: logp\n            costs += [-f for f in guide_terms[\"log_factors\"]]  # guide costs: -logq\n\n            # finally, integrate out guide variables in the elbo and all plates\n            plate_vars = guide_terms[\"plate_vars\"] | model_terms[\"plate_vars\"]\n            elbo = to_funsor(0, output=funsor.Real)\n            for cost in costs:\n                # compute the marginal logq in the guide corresponding to this cost term\n                log_prob = funsor.sum_product.sum_product(\n                    funsor.ops.logaddexp,\n                    funsor.ops.add,\n                    guide_terms[\"log_measures\"],\n                    plates=plate_vars,\n                    eliminate=(plate_vars | guide_terms[\"measure_vars\"])\n                    - frozenset(cost.inputs),\n                )\n                # compute the expected cost term E_q[logp] or E_q[-logq] using the marginal logq for q\n                elbo_term = funsor.Integrate(\n                    log_prob, cost, guide_terms[\"measure_vars\"] & frozenset(cost.inputs)\n                )\n                elbo += elbo_term.reduce(\n                    funsor.ops.add, plate_vars & frozenset(cost.inputs)\n                )\n\n        # evaluate the elbo, using memoize to share tensor computation where possible\n        with funsor.interpretations.memoize():\n            return -to_data(apply_optimizer(elbo))\n\n\n@copy_docs_from(_OrigTraceEnum_ELBO)\nclass TraceEnum_ELBO(ELBO):\n    def differentiable_loss(self, model, guide, *args, **kwargs):\n\n        # get batched, enumerated, to_funsor-ed traces from the guide and model\n        with plate(\n            size=self.num_particles\n        ) if self.num_particles > 1 else contextlib.ExitStack(), enum(\n            first_available_dim=(-self.max_plate_nesting - 1)\n            if self.max_plate_nesting\n            else None\n        ):\n            guide_tr = trace(guide).get_trace(*args, **kwargs)\n            model_tr = trace(replay(model, trace=guide_tr)).get_trace(*args, **kwargs)\n\n        # extract from traces all metadata that we will need to compute the elbo\n        guide_terms = terms_from_trace(guide_tr)\n        model_terms = terms_from_trace(model_tr)\n\n        # build up a lazy expression for the elbo\n        with funsor.terms.lazy:\n            # identify and contract out auxiliary variables in the model with partial_sum_product\n            contracted_factors, uncontracted_factors = [], []\n            for f in model_terms[\"log_factors\"]:\n                if model_terms[\"measure_vars\"].intersection(f.inputs):\n                    contracted_factors.append(f)\n                else:\n                    uncontracted_factors.append(f)\n            contracted_costs = []\n            # incorporate the effects of subsampling and handlers.scale through a common scale factor\n            for group_factors, group_vars in _partition(\n                model_terms[\"log_measures\"] + contracted_factors,\n                model_terms[\"measure_vars\"],\n            ):\n                group_factor_vars = frozenset().union(\n                    *[f.inputs for f in group_factors]\n                )\n                group_plates = model_terms[\"plate_vars\"] & group_factor_vars\n                outermost_plates = frozenset.intersection(\n                    *(frozenset(f.inputs) & group_plates for f in group_factors)\n                )\n                elim_plates = group_plates - outermost_plates\n                for f in funsor.sum_product.partial_sum_product(\n                    funsor.ops.logaddexp,\n                    funsor.ops.add,\n                    group_factors,\n                    plates=group_plates,\n                    eliminate=group_vars | elim_plates,\n                ):\n                    contracted_costs.append(model_terms[\"scale\"] * f)\n\n            # accumulate costs from model (logp) and guide (-logq)\n            costs = contracted_costs + uncontracted_factors  # model costs: logp\n            costs += [-f for f in guide_terms[\"log_factors\"]]  # guide costs: -logq\n\n            # compute expected cost\n            # Cf. pyro.infer.util.Dice.compute_expectation()\n            # https://github.com/pyro-ppl/pyro/blob/0.3.0/pyro/infer/util.py#L212\n            # TODO Replace this with funsor.Expectation\n            plate_vars = guide_terms[\"plate_vars\"] | model_terms[\"plate_vars\"]\n            # compute the marginal logq in the guide corresponding to each cost term\n            targets = dict()\n            for cost in costs:\n                input_vars = frozenset(cost.inputs)\n                if input_vars not in targets:\n                    targets[input_vars] = funsor.Tensor(\n                        funsor.ops.new_zeros(\n                            funsor.tensor.get_default_prototype(),\n                            tuple(v.size for v in cost.inputs.values()),\n                        ),\n                        cost.inputs,\n                        cost.dtype,\n                    )\n            with AdjointTape() as tape:\n                logzq = funsor.sum_product.sum_product(\n                    funsor.ops.logaddexp,\n                    funsor.ops.add,\n                    guide_terms[\"log_measures\"] + list(targets.values()),\n                    plates=plate_vars,\n                    eliminate=(plate_vars | guide_terms[\"measure_vars\"]),\n                )\n            marginals = tape.adjoint(\n                funsor.ops.logaddexp, funsor.ops.add, logzq, tuple(targets.values())\n            )\n            # finally, integrate out guide variables in the elbo and all plates\n            elbo = to_funsor(0, output=funsor.Real)\n            for cost in costs:\n                target = targets[frozenset(cost.inputs)]\n                logzq_local = marginals[target].reduce(\n                    funsor.ops.logaddexp, frozenset(cost.inputs) - plate_vars\n                )\n                log_prob = marginals[target] - logzq_local\n                elbo_term = funsor.Integrate(\n                    log_prob,\n                    cost,\n                    guide_terms[\"measure_vars\"] & frozenset(log_prob.inputs),\n                )\n                elbo += elbo_term.reduce(\n                    funsor.ops.add, plate_vars & frozenset(cost.inputs)\n                )\n\n        # evaluate the elbo, using memoize to share tensor computation where possible\n        with funsor.interpretations.memoize():\n            return -to_data(apply_optimizer(elbo))\n\n\nclass JitTraceEnum_ELBO(Jit_ELBO, TraceEnum_ELBO):\n    pass\n\n\nclass JitTraceMarkovEnum_ELBO(Jit_ELBO, TraceMarkovEnum_ELBO):\n    pass\n", "path": "pyro/contrib/funsor/infer/traceenum_elbo.py"}]}