State-space model: is lax.scan compatible with numpyro.sample?

numpyro
8

Hi @neerajprad. Yes, this makes complete sense, thanks! I had considered a solution like this where the key is passed around as part of the state:

def transition(state, i, q=1., r=1., phi=0.5, beta=0.5):
    x0, mu0, key = state
    key, subkey1, subkey2 = jax.random.split(key, 3)
    x1 = numpyro.sample(f'x_{i}', dist.Normal(phi*x0, q), rng_key=subkey1)
    mu1 = beta * mu0 + x1
    y1 = numpyro.sample(f'y_{i}', dist.Normal(mu1, r), rng_key=subkey2)
    return (x1, mu1, key), (x1, y1)

but I didn’t know how to connect this kind of explicit key handling to the numpyro handlers (i.e. PRNGIdentity() distribution and the rng_key argument to numpyro.sample). This works perfectly, thanks!