Affine Coupling Quick Question

Hi!

I’m implementing Normalizing Flows for zero-shot learning.

I have a quick question about https://docs.pyro.ai/en/stable/distributions.html#pyro.distributions.transforms.AffineCoupling. What is the purpose of param_dims in defining the hypernet?

Thanks!

cc @stefanwebb

param_dims defines the output shape of the conditioning network. For the AffineCoupling transform it has a single mean and log(variance) for each dimension that is modelled, and there are input_dim-split_dim dimensions