I want to use an inverse autoregressive flow (IAF) for the guide distribution in a Bayesian neural network. I'm not sure how to specify the IAF distribution for the weight matrices because these are 2D. Previously I was using a
Normal distribution, and I can just create a location and scale with the same shape as the parameter (e.g. 512 x 100). However,
InverseAutoregressiveFlow only accepts an
int for the input dimension. I assume I'll need to have
input_dim=512 * 100, but then the samples won't be the right shape, so I expect there will be an error in SVI (or will it reshape for me?). Should I define a new
Transform that does the reshaping? Or is there a better way to do this?