Subsampling in autoguide / easyguide (?)

Hello there,
My model contains a loop of the form:

for t in pyro.markov(torch.arange(T)):
    pyro.sample("z_{}".format(t), ...)

In order for the guide to be fast, I would like it to generate iid z_t, but I haven’t found a way to put these in a vectorized plate, since their names are different. A possible way out is an automatic guide such as AutoDiagonalNormal, which seems to perform well when generating all the z_t.

However, I’m not sure how plate subsampling fits into the picture. According to https://github.com/pyro-ppl/pyro/pull/1888, it is indeed missing from pyro.contrib.autoguide but feasible in pyro.contrib.easyguide (http://docs.pyro.ai/en/dev/contrib.easyguide.html). Could someone provide an example of this? Maybe @fritzo ?

Thanks in advance
Giom

Check out the easyguide tutorial. Alternatively, can you relax your requirement that all zs have different names? I’ve had success combining torch.stack with partial vectorization, e.g.

def partially_vectorized_model(data):
    # Sample non-vectorized part in a loop.
    xs = [None] * len(data)
    for t in pyro.markov(range(1, len(data))):
        x_prev = xs[t - 1] if t else 0.
        xs[t] = pyro.sample("x_{}".format(t), dist.Normal(x_prev, 1.))

    # Stack and sample remaining sites in a plate.
    xs = torch.stack(xs)
    with pyro.plate("data", len(data)):
        z = pyro.sample("z", dist.Normal(xs, 1.))
        pyro.sample("obs", dist.Normal(z, 1.), obs=data)

Thanks for the answer! Since I don’t fully understand automatically-generated guides, I’ve stuck with the “auxiliary-unpacked-to-Delta” method suggested in the tutorial, which already speeds my code up by a factor of 2.