Using guide for downstream sample sites of enumerated site - Issues

Hey guys,

As far as I can tell, pyro does not allow a sample site used in a guide to be downstream of a discrete sample site, when the discrete sample site is using parallel inference. For example, if I have the simple structure:

for t, y in pyro.markov(enumerate(data)):
      s = pyro.sample(f"s_{t}", dist.Categorical(transition_matrix[s]),infer={"enumerate": "parallel"})

      mask = s==1

      dist_z = dist.Normal(
                      mask * z * Vindex(x_ar_coeff)[1] 
                      + (~mask) * z * Vindex(x_ar_coeff)[0]

      z = pyro.sample(f"z_{t}", dist_z)


, and even if I use an autoguide, Pyro will throw an error saying that the guide and model sample sites do not match at site z_{t}, for any t. Obviously the reason is that during enumeration, dist_z will be expanded to the shape of the fully-enumerated variable s, but Pyro does not seem to allow this, since the guide and model shapes have to match. Is there a way to get the sample shapes to match at the sample site pyro.sample(f"z_{t}", dist_z)? Broadcasting from the guide to the model would fix the problem, but pyro doesn’t broadcast here, instead throwing an error.

Hi @mike_schoehals , can you provide a minimal reproducible code (with all imports) please? I can have a look at it.


Thanks for taking a look at. Here is a full end-to-end example:

import numpy as np
import torch
from torch.distributions import constraints
import pyro
from pyro import poutine
import pyro.distributions as dist
from pyro.infer import SVI, Trace_ELBO, TraceEnum_ELBO, TraceGraph_ELBO 
from pyro.infer.autoguide import AutoNormal, AutoDiagonalNormal

N = 10
data = torch.randn(N)


def model(data,is_enum=False):
    trans_prob = pyro.param("trans_prob",torch.randn(2).exp(),constraint=constraints.simplex)
    transition_matrix = torch.stack([trans_prob,torch.flip(trans_prob,(0,))],dim=0)


    for t, y in pyro.markov(enumerate(data)):
        #for t, y in enumerate(data):
        s = pyro.sample(f"s_{t}", dist.Categorical(transition_matrix[s]),infer={"enumerate": "parallel"})

        mask = s==1

        dist_z = dist.Normal(
                        mask * z * 0.99
                        + (~mask) * z * 0.5
        z = pyro.sample(f"z_{t}", dist_z)

        obs = pyro.sample(f"y_{t}", dist.Normal(z,1).to_event(1), obs=y)

        # this is required in this regime-switching state-space model.
        # z will be constant along the enum dimension, and we only want 1 value to continue to the next iteration of the loop.
        if is_enum:
            z = z[0,...].squeeze().unsqueeze(-1)

optim = pyro.optim.Adam({'lr': 1e-1})
guide =  AutoDiagonalNormal(poutine.block(model,hide=[f's_{t}' for t in range(data.shape[0])]))

svi = SVI(model,guide, optim, loss=TraceEnum_ELBO(retain_graph=True,max_plate_nesting=1))#,num_particles=10,vectorize_particles=True))
for i in range(2000):
    loss = svi.step(data,True)
    if not i % 10:
        print('loss: ', loss)

@ordabayev No worries if you haven’t looked at this, but if you had, I was curious if you’d found anything. This case is useful for regime-switching time series, and being able to use parallel enumeration allows for very low variance estimates (as you guys are aware).

Hi @mike_schoehals , sorry for not replying. I had an initial look and couldn’t figure it out. Didn’t have time since then to have a deeper look. I’ll post it here if I have any progress.

@ordabayev no worries. thanks for the reply!