add comment for why we divide by subsample_size at z_0 log prob

pyro-ppl · Jun 30, 2023 · 3450771 · 3450771
1 parent 96a7be4
commit 3450771
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
@@ -1420,8 +1420,11 @@ def base_z_dist_log_prob(z):
                                 / scale
                             )
 
+                # The log_prob of z_0 will be broadcasted to `subsample_size` because this statement
+                # is run under the subsample plate. Hence we divide the log_prob by `subsample_size`.
                 numpyro.factor(
-                    "{}_z_0_factor".format(self.prefix), base_z_dist_log_prob(z_0) / subsample_size
+                    "{}_z_0_factor".format(self.prefix),
+                    base_z_dist_log_prob(z_0) / subsample_size,
                 )
             else:
                 z_0_loc_init = jnp.zeros((N, D))