Parameterizing normalizing flows over model parameters

Hi all,
I have a hierarchical model that is parametrized by parameter a.
From this model, I can generate samples for a given a:
x,y ~ P(x, y | a=a*)

I would now like to parameterize the distribution P(x, y | a) using normalizing flows. For a given a=a* and the corresponding data points x, y I can parametrize the PDF nicely - but how do I include the dependency on the model parameter a? I assume I would have to add an additional input (for a) to the MLP’s that learn the flow parameters?

Here’s a toy model that has the same features as my real example:

a = 1
# sample x ~ P(x | a= 1)
x_a_1 = scipy.stats.norm(a).rvs(1000)
dist = fit_flow(x_a_1)

a = 2
# sample x ~ P(x | a= 2)
x_a_2 = scipy.stats.norm(a).rvs(1000)
dist2 = fit_flow(x_a_2)

please see this part of the normalizing flow tutorial

Thanks, I might be able to make that work for my case.
Technical question regarding the example: why is x1.detach()and x2.detach() used in the evaluation of the conditional pdf?