extending Predictive to support sampling by plate name

Following our earlier discussion with @fritzo and @adamgayoso , creating this issue. Would be great if Predictive class can support sampling of all variables in a given plate, automatically pulling all site values in a given plate(plate_name) or all sites not is that plate, including both samples and observed values:

trace = poutine.trace(model).get_trace(*args, **kwargs)  # you'll need to fill in args,kwargs
cell_samples = {
    name: site["value"]
    for name, site in trace.nodes.items()
    if site["type"] == "sample"
    if any(f.name == plate_name for f in site["cond_indep_stack"])
}

This functionality would simplify dealing with local variables within a minibatch plate, by making it easy to do separate predictive calls for the minibatch plate variables and global variables.

[vitkl]

Yes, exactly like @jamestwebber described - we need a way to sample both the global variables and the local variables (including amortised variables). The idea is that by doing predictive sampling by plate name you can isolate local variables from global variables and sample them via 2 separate calls. I think 2 separate calls are needed because if you iterate over the data loader global variables will be sampled in every iteration of the data loader whereas local variable will be only sampled once in their respective minibatch. Separating allows sampling global variable once (where I count y_fg as global: d_cg = x_cf @ y_fg and x_cf is minibatched in c plate).

I see 2 separate problems:

1. How to tell predictive to sample local variables from a batch plate
2. How to concatenate the batches of local variables correctly (which need recognising the plate dimension)

[jamestwebber]

it sounds like Predictive is missing support for vectorization somewhere, but I suspect I’m misunderstanding.

It’s very possible that I just don’t understand the right pattern for accomplishing this. For me, the intended workflow is something like

1. Learn a model of some large dataset using mini-batches for training
2. Make a Predictive instance with pred = Predictive(model, guide=guide, num_samples=1000)
3. Collect the posterior on my large dataset by iterating through the same data loader

posteriors = []
for x, y in data_loader:
    posteriors.append(pred(x, y))

aggregate_posterior_minibatches()

So that in the end I’d have a posterior over both the global variables (i.e. some background effect) and the minibatch-specific ones (for example, cell-specific loading factors).