Struggling to find the code that correspond to Stochastic Inverse

[Donglin-Wang2]

I was trying to read the source code and match it to the original paper. However, I was struggling to find the code for the stochastic inverse part of the propagation. From my understanding, the forward direction of the implementation is the exact opposite of the paper. Namely, in the code, $f_\theta(x)$ is mapping from X to Z, which opposite from the paper. For code tracing, I try to see which part of the code would generate Fig 4 in the paper (1st screenshot below) following the implementation of class SimpleNIF(GLOW) in models.py. I managed to trace the code from the def save_test_embeddings(key, experiment, save_path, n_samples_per_batch=64) in the evaluate_experiments.py to the def sequential_flow(*layers) method in the normalizing_flows.py file. It seems to me that this method is only taking the direct inverse of Glow without applying the stochastic inverse formula $q_{\theta}(z \mid x)=\frac{p_{\theta}(x \mid z)}{\int p_{\theta}\left(x \mid z^{\prime}\right) d z^{\prime}}$ in Eq. 7 of the paper (2nd screenshot below). Can you point out a piece of code that directly corresponds to Eq. 7 in the paper?

[EddieCunningham]

Thanks for the interest in this project! There are actually a few implementations of equation 7 and all of that code for that is in this file. There are implementations for using a diagonal/full covariance, closed form prior / stochastic inverse, coupling/no coupling and nearest neighbors upsampling. The SimpleNIF model uses the closed form prior (AffineGaussianPriorDiagCov) while the NIF model uses the stochastic inverse (TallAffineDiagCov) and a lower dimensional flow over z.

[Donglin-Wang2]

When I was reading the analytical solution for the Diagonal Covariance case for Gaussian inverse, I have some questions regarding the notation. In the picture below, I am wondering what the N^(-1) means in equation 41 and how you went about going from 40 to 41.

[EddieCunningham]

N^(-1) corresponds to the information form of a multivariate gaussian. We went from step 40 to 41 by recognizing that the first 2 terms of equation 40 are proportional to the density under an information form gaussian and we then multiplied in the terms necessary to make sure equation 40 and 41 are the same.