fixed log_prob bug in snre

CHANGELOG.md

@@ -1,3 +1,7 @@
+# v0.11.1
+- Bug fix for log_prob() in SNRE (#275)
+
+
 # v0.11.0
 
 - Changed the API to do multi-round inference (#273)

sbi/inference/posterior.py

@@ -317,23 +317,24 @@ def _log_prob_snpe(self, theta: Tensor, x: Tensor, norm_posterior: bool) -> Tens
         return masked_log_prob - log_factor
 
     def _log_prob_ratio_estimator(self, theta: Tensor, x: Tensor) -> Tensor:
-        log_ratio = self.net(torch.cat((theta, x)).reshape(1, -1))
+        log_ratio = self.net(torch.cat((theta, x), dim=1)).reshape(-1)
         return log_ratio + self._prior.log_prob(theta)
 
     def _log_prob_snre_a(self, theta: Tensor, x: Tensor) -> Tensor:
-        warn(
-            "The log probability from SRE is only correct up to a normalizing constant."
-        )
-        return self._log_prob_ratio_estimator(theta, x)
-
-    def _log_prob_snre_b(self, theta: Tensor, x: Tensor) -> Tensor:
         if self._num_trained_rounds > 1:
             warn(
-                "The log-probability from AALR beyond round 1 is only correct "
+                "The log-probability from AALR / SNRE-A beyond round 1 is only correct "
                 "up to a normalizing constant."
             )
         return self._log_prob_ratio_estimator(theta, x)
 
+    def _log_prob_snre_b(self, theta: Tensor, x: Tensor) -> Tensor:
+        warn(
+            "The log probability from SNRE_B is only correct up to a normalizing "
+            "constant."
+        )
+        return self._log_prob_ratio_estimator(theta, x)
+
     def _log_prob_snle(self, theta: Tensor, x: Tensor) -> Tensor:
         warn(
             "The log probability from SNL is only correct up to a normalizing constant."

tests/linearGaussian_snpe_test.py

@@ -83,7 +83,7 @@ def simulator(theta):
             posterior, x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov
         )
 
-        max_dkl = 0.1 if num_dim == 1 else 0.8
+        max_dkl = 0.1
 
         assert (
             dkl < max_dkl

tests/linearGaussian_snre_test.py

@@ -231,7 +231,7 @@ def simulator(theta):
             posterior, x_o[0], likelihood_shift, likelihood_cov, prior_mean, prior_cov
         )
 
-        max_dkl = 0.1 if num_dim == 1 else 0.8
+        max_dkl = 0.1
 
         assert (
             dkl < max_dkl