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

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)

[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?