fix: posterior potential iid handling (#1276)

Author: janfb

sbi/inference/potentials/posterior_based_potential.py

@@ -89,11 +89,11 @@ def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = False):
         For posterior-based methods, `x_o` is not allowed to be iid, as we assume that
         iid `x` is handled by a Permutation Invariant embedding net.
         """
-        if x_is_iid:
+        if x_is_iid and x_o is not None and x_o.shape[0] > 1:
             raise NotImplementedError(
-                "For NPE, iid `x` must be handled by a Permutation Invariant embedding \
-                    net. Therefore, the iid dimension of `x` is added to the event\
-                        dimension of `x`. Please set `x_is_iid=False`."
+                "For NPE, iid `x` must be handled by a permutation invariant embedding "
+                "net. Therefore, the iid dimension of `x` is added to the event "
+                "dimension of `x`. Please set `x_is_iid=False`."
             )
         else:
             super().set_x(x_o, x_is_iid=False)

tests/linearGaussian_fmpe_test.py

@@ -287,8 +287,8 @@ def simulator(theta):
     # Evaluate the conditional density be drawing samples and smoothing with a Gaussian
     # kde.
     potential_fn, theta_transform = posterior_estimator_based_potential(
-        posterior_estimator, prior=prior, x_o=x_o
-    )
+        posterior_estimator, prior=prior
+    ).set_x(x_o, x_is_iid=False)
     (
         conditioned_potential_fn,
         restricted_tf,

tests/linearGaussian_snpe_test.py

@@ -558,8 +558,8 @@ def simulator(theta):
     # Evaluate the conditional density be drawing samples and smoothing with a Gaussian
     # kde.
     potential_fn, theta_transform = posterior_estimator_based_potential(
-        posterior_estimator, prior=prior, x_o=x_o
-    )
+        posterior_estimator, prior=prior
+    ).set_x(x_o, x_is_iid=False)
     (
         conditioned_potential_fn,
         restricted_tf,