Fix iid error handling

sbi/inference/potentials/likelihood_based_potential.py

@@ -60,7 +60,7 @@ def __init__(
         self,
         likelihood_estimator: ConditionalDensityEstimator,
         prior: Distribution,
-        x_o: Optional[Tensor],
+        x_o: Optional[Tensor] = None,
         device: str = "cpu",
     ):
         r"""Returns the potential function for likelihood-based methods.

sbi/inference/potentials/posterior_based_potential.py

@@ -63,7 +63,7 @@ def __init__(
         self,
         posterior_estimator: ConditionalDensityEstimator,
         prior: Distribution,
-        x_o: Optional[Tensor],
+        x_o: Optional[Tensor] = None,
         device: str = "cpu",
     ):
         r"""Returns the potential for posterior-based methods.
@@ -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)

sbi/inference/potentials/ratio_based_potential.py

@@ -51,7 +51,7 @@ def __init__(
         self,
         ratio_estimator: nn.Module,
         prior: Distribution,
-        x_o: Optional[Tensor],
+        x_o: Optional[Tensor] = None,
         device: str = "cpu",
     ):
         r"""Returns the potential for ratio-based methods.

sbi/inference/potentials/score_based_potential.py

@@ -60,7 +60,7 @@ def __init__(
         self,
         score_estimator: ConditionalScoreEstimator,
         prior: Optional[Distribution],
-        x_o: Optional[Tensor],
+        x_o: Optional[Tensor] = None,
         iid_method: str = "iid_bridge",
         device: str = "cpu",
     ):

tests/linearGaussian_fmpe_test.py

@@ -300,6 +300,7 @@ def simulator(theta):
         condition=samples[0],
         dims_to_sample=[dim_to_sample_1, dim_to_sample_2],
     )
+    conditioned_potential_fn.set_x(x_o, x_is_iid=False)
     mcmc_posterior = MCMCPosterior(
         potential_fn=conditioned_potential_fn,
         theta_transform=restricted_tf,

tests/linearGaussian_snpe_test.py

@@ -571,6 +571,7 @@ def simulator(theta):
         condition=samples[0],
         dims_to_sample=[dim_to_sample_1, dim_to_sample_2],
     )
+    conditioned_potential_fn.set_x(x_o, x_is_iid=False)
     mcmc_posterior = MCMCPosterior(
         potential_fn=conditioned_potential_fn,
         theta_transform=restricted_tf,

tutorials/05_conditional_distributions.ipynb

@@ -396,12 +396,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
+    "x_o = torch.ones(1, 4)\n",
     "potential_fn, theta_transform = posterior_estimator_based_potential(\n",
-    "    posterior_estimator, prior=prior, x_o=torch.ones(4)\n",
+    "    posterior_estimator, prior=prior, x_o=x_o\n",
     ")"
    ]
   },