Hi, great work creating these ML models. I have recreated a subset of the TotalVI VAE model handling ADT data in Tensorflow. I the process I noticed that the distribution for the background rate is a function of the z latent variable. However I can not find any documentation for that in the TotalVI paper. There it is only stated that the prior (beta_nt) for that distribution is Lognormal with mean/sd for each batch X protein, but I don’t see any reference stating that the posterior is computed from z (while the functions computing the foreground scale and the mixing are explicitly described) So, should the background rate depend on Z, or should it only depend on the batch?