Repository: pyro-ppl/pyro

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 


Discussion:
I suspect this is due to a torch==1.11.0 issue with `ProvenanceTensor` which #3045 fixed via

<details>

```diff
diff --git a/pyro/ops/provenance.py b/pyro/ops/provenance.py
index b1ec32d9..51a1ec9e 100644
--- a/pyro/ops/provenance.py
+++ b/pyro/ops/provenance.py
@@ -46,9 +46,7 @@ class ProvenanceTensor(torch.Tensor):
         assert not isinstance(data, ProvenanceTensor)
         if not provenance:
             return data
-        instance = torch.Tensor.__new__(cls)
-        instance.__init__(data, provenance)
-        return instance
+        return super().__new__(cls)

     def __init__(self, data, provenance=frozenset()):
         assert isinstance(provenance, frozenset)
```

</details>

and I suspect we could apply the same fix to funsor (and release, and pin Pyro to that version).

EDIT I've implemented this idea in the simple PR https://github.com/pyro-ppl/funsor/pull/588

EDIT this didn't seem to be the cause 🤷 
Here's some [jamboarding](https://jamboard.google.com/d/1lc0gpUPMrU2M0QMn2cSUH1M6Pyct-l2OzRMvX8wfwoo) @ordabayevy and I did to examine models.

We currently believe the bug may actually be in `pyro.contrib.funsor.infer.TraceEnum_ELBO` in that some likelihood-only plates are not being passed to `eliminate`.

Resolving patch:
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 @@ def differentiable_loss(self, model, guide, *args, **kwargs):
                     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


Test patch:
diff --git a/tests/contrib/funsor/test_enum_funsor.py b/tests/contrib/funsor/test_enum_funsor.py
--- a/tests/contrib/funsor/test_enum_funsor.py
+++ b/tests/contrib/funsor/test_enum_funsor.py
@@ -264,7 +264,6 @@ def guide():
     _check_loss_and_grads(hand_loss, auto_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pytest.mark.parametrize(
     "num_samples,num_masked", [(2, 2), (3, 2)], ids=["batch", "masked"]
@@ -499,7 +498,6 @@ def hand_guide(data):
     _check_loss_and_grads(hand_loss, auto_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pytest.mark.parametrize(
     "outer_obs,inner_obs", [(False, True), (True, False), (True, True)]
@@ -642,7 +640,6 @@ def guide_iplate():
         _check_loss_and_grads(expected_loss, actual_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("enumerate1", ["parallel", "sequential"])
 @pyroapi.pyro_backend(_PYRO_BACKEND)
 def test_elbo_enumerate_plate_6(enumerate1):
@@ -705,7 +702,6 @@ def guide():
     _check_loss_and_grads(expected_loss, actual_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pyroapi.pyro_backend(_PYRO_BACKEND)
 def test_elbo_enumerate_plate_7(scale):
@@ -877,7 +873,6 @@ def guide():
     _check_loss_and_grads(hand_loss, auto_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pyroapi.pyro_backend(_PYRO_BACKEND)
 def test_elbo_enumerate_plates_2(scale):
@@ -934,7 +929,6 @@ def guide():
     _check_loss_and_grads(hand_loss, auto_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pyroapi.pyro_backend(_PYRO_BACKEND)
 def test_elbo_enumerate_plates_3(scale):
@@ -987,7 +981,6 @@ def guide():
     _check_loss_and_grads(hand_loss, auto_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pyroapi.pyro_backend(_PYRO_BACKEND)
 def test_elbo_enumerate_plates_4(scale):
@@ -1047,7 +1040,6 @@ def guide(data):
     _check_loss_and_grads(hand_loss, auto_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pyroapi.pyro_backend(_PYRO_BACKEND)
 def test_elbo_enumerate_plates_5(scale):
@@ -1111,7 +1103,6 @@ def guide():
     _check_loss_and_grads(hand_loss, auto_loss)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pyroapi.pyro_backend(_PYRO_BACKEND)
 def test_elbo_enumerate_plates_6(scale):
@@ -1250,7 +1241,6 @@ def guide(data):
         elbo.differentiable_loss(model_plate_plate, guide, data)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("scale", [1, 10])
 @pyroapi.pyro_backend(_PYRO_BACKEND)
 def test_elbo_enumerate_plates_7(scale):
@@ -1403,7 +1393,6 @@ def guide(data):
     _check_loss_and_grads(loss_iplate_iplate, loss_plate_plate)
 
 
-@pytest.mark.xfail(reason="https://github.com/pyro-ppl/pyro/issues/3046")
 @pytest.mark.parametrize("guide_scale", [1])
 @pytest.mark.parametrize("model_scale", [1])
 @pytest.mark.parametrize(