pyro-ppl/pyro

extending Predictive to support sampling by plate name

Open

#2,783 opened on Mar 16, 2021

View on GitHub
 (5 comments) (1 reaction) (0 assignees)Python (981 forks)batch import
discussionenhancementhelp wanted

Repository metrics

Stars
 (8,211 stars)
PR merge metrics
 (Avg merge 10d 19h) (1 merged PR in 30d)

Description

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.

Contributor guide