Is it possible to have a normalizing flow as a latent parameter in a model in Pyro? For example, if I’m trying to model an effect that is spatially-varying, is it possible to express this as a normalizing flow? For example, schematically:
NF = transformed_dist(N([0,0], [1,1]), spline)
beta ~ NF
y = N(mu + beta, sigma)
Or can the NF only be used for observed variables?
the NF is just a fancy parametric distribution so it can be used just like any other distribution.
That’s what I had hoped, but its not clear how this works. If the NF is just like another distribution, I would expect to do something like this:
mu = pyro.sample('mu', Normal(loc=torch.tensor([0.]), scale=torch.tensor([1.])))
base_dist = Normal(torch.zeros(2), torch.ones(2))
spline_transform = spline_coupling(2, count_bins=16)
theta = pyro.sample('theta', TransformedDistribution(base_dist, [spline_transform]))
loc = mu + theta
with pyro.plate("obs", len(dataset)):
pyro.sample("y", Normal(loc=loc, scale=torch.tensor(1.)), obs=torch.tensor(y))
But clearly this does not work. Do I need to pass input coordinates to theta somehow?
can you please provide error messages and/or complete code snippets that are runnable?
Never mind, I think I am misapplying normalizing flows here. Apologies.