Skip to content

Support irange/iarange nested inside iarange #370

@fritzo

Description

@fritzo

See Design Doc

This proposes generalizing TraceGraph_ELBO to apply variance reduction techniques when an irange or iarange is nested beneath an iarange.

Now that Pyro distributions support arbitrarily many batch dimensions, TraceGraph_ELBO should take care to sum out only those dimensions of each sample site that are controlled by an iarange. For example, consider this model (say with a mean field guide)

def model(time_series):
    N = time_series.size(-1)
    with iarange("series", N):
        z = pyro.sample("z", Bernoulli(0.5 * ng_ones(365, N)))
        pyro.sample("obs", Bernoulli(0.5 * ng_ones(365, N)),
                    obs=(time_series * x))

We need to ensure that TraceGraph_ELBO sums out dimension 0 of size 365 but does not sum out dimension 1 of size N (since that dimension is declared independent by the iarange).

Metadata

Metadata

Type

No type

Projects

No projects

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions