pyro.render_model() fails to draw overlapping non-nested plates

[fritzo]

Due to a shortcoming in graphviz clusters, Pyro's and NumPyro's pyro.render_model() currently cannot render overlapping non-nested plates, e.g. the Venn diagram

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)

[eb8680]

One workaround would be to add an option to render with daft, which handles plates correctly (example).

[fritzo]

One optional would be to replace graphviz by daft-pgm, which does support nested plates. However daft-pgm does not support automatic node layout.

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?

[fehiepsi]

@kpj mentioned that we can use graphviz to find positions then use daft for plotting. Probably that will work for a class of models.

[fehiepsi]

It seems that in numpyro, we have

which can be seen as a temporary solution. What do you think @fritzo ?

[fritzo]

@fehiepsi NumPyro's current behavior looks great, and I suspect it's the only fully general solution 👍 Do you know if that was a recent change, or did I make an error when porting code to Pyro #2962?