Latent normalizing flows?

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.