Repository: pyro-ppl/pyro

Issue:
Possible bug in poutine/trace_messenger.py, function identify_dense_edges
Hi,

I found some strange code pattern in function identify_dense_edges, in file [poutine/trace_messenger.py](https://github.com/pyro-ppl/pyro/blob/dev/pyro/poutine/trace_messenger.py#L21).
I notice two sequential instances of:
`if site_is_subsample(node):
            continue`
The first one appears at line 17 (in the main loop of the function) and the second one at line 21 (in the nested loop).
It seems the second condition will never be reached; if it is, the first one would be selected, going to the next iteration of the main loop.

Should the second not be as follows ?
`if site_is_subsample(past_node):
            continue`

(this would be consistent with considering only pairs of non subsample nodes).

Thanks,

Regards,

Xavier


Discussion:
Hi @xrival, indeed that line looks like a bug. Feel free to submit a one-line fix PR, or wait until we get around to fixing it ourselves. Many eyes 👀 
Hi @fritzo,
Many thanks! I will try to do so in the next few days.
Hi @xrival mind if I quick fix this to make it into our 1.8.6 release?  Thanks again for reporting it.

Resolving patch:
diff --git a/pyro/poutine/trace_messenger.py b/pyro/poutine/trace_messenger.py
--- a/pyro/poutine/trace_messenger.py
+++ b/pyro/poutine/trace_messenger.py
@@ -18,7 +18,7 @@ def identify_dense_edges(trace):
             continue
         if node["type"] == "sample":
             for past_name, past_node in trace.nodes.items():
-                if site_is_subsample(node):
+                if site_is_subsample(past_node):
                     continue
                 if past_node["type"] == "sample":
                     if past_name == name:


Test patch:
diff --git a/tests/integration_tests/test_tracegraph_elbo.py b/tests/integration_tests/test_tracegraph_elbo.py
--- a/tests/integration_tests/test_tracegraph_elbo.py
+++ b/tests/integration_tests/test_tracegraph_elbo.py
@@ -472,7 +472,7 @@ def guide():
         model_trace = pyro.poutine.trace(
             pyro.poutine.replay(model, trace=guide_trace), graph_type="dense"
         ).get_trace()
-        assert len(list(model_trace.edges)) == 27
+        assert len(list(model_trace.edges)) == 9
         assert len(model_trace.nodes) == 16
         assert len(list(guide_trace.edges)) == 0
         assert len(guide_trace.nodes) == 9
diff --git a/tests/poutine/test_poutines.py b/tests/poutine/test_poutines.py
--- a/tests/poutine/test_poutines.py
+++ b/tests/poutine/test_poutines.py
@@ -846,6 +846,20 @@ def model():
     assert len(_DIM_ALLOCATOR._stack) == 0, "stack was not cleaned on error"
 
 
+@pytest.mark.parametrize(
+    "graph_type, expected", [("flat", set()), ("dense", {"x", "y"})]
+)
+def test_trace_plate(graph_type: str, expected: set):
+    def model():
+        with pyro.plate("plate", 2):
+            x = pyro.sample("x", dist.Normal(0, 1))
+            pyro.sample("y", dist.Normal(x, 1))
+
+    trace = poutine.trace(model, graph_type=graph_type).get_trace()
+    nodes = set().union(*trace._succ.values(), *trace._pred.values())
+    assert nodes == expected
+
+
 def test_decorator_interface_primitives():
     @poutine.trace
     def model():