Condition IAF flows on context input h?

I was looking at the normalizing flows module as I am hoping to use them on a project I’ve been looking at. Is it possible to pass additional inputs to the IAF layers before creating a TransformedDistribution? I was re-reading the IAF paper and they mention passing as input a context vector h to each flow besides the previous samples z_{t-1}.

From looking at the Pyro code for it, it doesn’t seem they support giving a context h to each flow. Is this actually the case, or did I just miss something? If it ISN"T supported, is there any suggested practices I should be aware of to make an IAF flow that does allow this?

So maybe something like:

#necessary imports (Gaussian, torch.nn , etc. etc.)
iaf = IAFCondition(...)
mu_0 eps_0, h = nn(X)
iaf.condition(h)
dist = Gaussian(mu_0, sigma_0)
dist = TransformedDistribution(dist, [iaf])
#the rest of my pyro goodness