Repository: pyro-ppl/pyro

Issue:
pyro.render_model() fails to draw overlapping non-nested plates
Due to [a shortcoming](https://stackoverflow.com/questions/12680865/how-to-draw-overlapping-clusters-in-graphviz) in [graphviz clusters](https://graphviz.org/Gallery/directed/cluster.html), Pyro's and NumPyro's [pyro.render_model()](https://docs.pyro.ai/en/latest/infer.util.html#pyro.infer.inspect.render_model) currently cannot render overlapping non-nested plates, e.g. the Venn diagram
```p
def model():
    plate1 = pyro.plate("plate1", 2, dim=-2)
    plate2 = pyro.plate("plate2", 3, dim=-1)
    with plate1:
        x = pyro.sample("x", dist.Normal(0, 1))
    with plate1, plate2:
        y = pyro.sample("y", dist.Normal(x, 1))
    with plate2:
        pyro.sample("z", dist.Normal(y.sum(-2, True), 1), obs=torch.zeros(3))

pyro.render_model(model)
```
![image](https://user-images.githubusercontent.com/648532/145103608-b0dcf75f-8a77-4b56-88dc-25d787a17b0b.png)


Discussion:
One workaround would be to add an option to render with [daft](https://docs.daft-pgm.org/en/latest/), which handles plates correctly ([example](https://docs.daft-pgm.org/en/latest/examples/deconvolution/)).
One optional would be to replace `graphviz` by [daft-pgm](https://docs.daft-pgm.org/en/latest/), which does support nested plates. However daft-pgm [does not support automatic node layout](https://github.com/daft-dev/daft/issues/122).

A more basic problem is that not all plate nesting patterns can be drawn by rectangular overlapping plates. Maybe we could use a generic venn diagram plotting library?
@kpj [mentioned](https://github.com/pyro-ppl/numpyro/issues/949#issuecomment-793549236) that we can use graphviz to find positions then use daft for plotting. Probably that will work for a class of models.
It seems that in numpyro, we have
![image](https://user-images.githubusercontent.com/4736342/145137100-16375ca4-6dfb-4eb9-a737-f7f76d27df03.png)
which can be seen as a temporary solution. What do you think @fritzo ?
@fehiepsi NumPyro's current behavior looks great, and I suspect it's the only fully general solution :+1: Do you know if that was a recent change, or did I make an error when porting code to Pyro #2962?

Resolving patch:
diff --git a/pyro/infer/inspect.py b/pyro/infer/inspect.py
--- a/pyro/infer/inspect.py
+++ b/pyro/infer/inspect.py
@@ -293,6 +293,24 @@ def model(data):
         if site["is_observed"]:
             observed.append(name)
 
+    def _resolve_plate_samples(plate_samples):
+        for p, pv in plate_samples.items():
+            pv = set(pv)
+            for q, qv in plate_samples.items():
+                qv = set(qv)
+                if len(pv & qv) > 0 and len(pv - qv) > 0 and len(qv - pv) > 0:
+                    plate_samples_ = plate_samples.copy()
+                    plate_samples_[q] = pv & qv
+                    plate_samples_[q + "__CLONE"] = qv - pv
+                    return _resolve_plate_samples(plate_samples_)
+        return plate_samples
+
+    plate_sample = _resolve_plate_samples(plate_sample)
+    # convert set to list to keep order of variables
+    plate_sample = {
+        k: [name for name in trace.nodes if name in v] for k, v in plate_sample.items()
+    }
+
     return {
         "sample_sample": sample_sample,
         "sample_dist": sample_dist,


Test patch: