Repository: sbi-dev/sbi

Issue:
iid samples support via NeuralPosteriorEnsemble (NLE as a base)
Hi,

When using `NeuralPosteriorEnsemble` with NLE as a base and multiple iid observations I get a value error stating that "The x passed to condition the posterior for evaluation or sampling has an inferred batch shape larger than one."

Here is a minimal code example that replicates the error:
```python
import torch
from sbi import analysis as analysis
from sbi import utils as utils
from sbi.inference import (SNLE, MCMCPosterior,
                           likelihood_estimator_based_potential,
                           prepare_for_sbi, simulate_for_sbi)
from sbi.utils.posterior_ensemble import NeuralPosteriorEnsemble

num_dim = 3
prior = utils.BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))

en_x_o = 10
x_o = torch.zeros(len_x_o, num_dim,)+ 0.01*torch.randn(len_x_o, num_dim,)

def linear_gaussian(theta):
    return theta + 1.0 + torch.randn_like(theta) * 0.1

def get_posterior(x_o, prior, simulator, num_rounds=2, num_simulations=3000, num_chains=40):
    simulator, prior = prepare_for_sbi(simulator, prior)
    inference = SNLE(prior=prior, density_estimator="nsf")

    proposal = prior
    for _ in range(num_rounds):
        theta, x = simulate_for_sbi(simulator, proposal, num_simulations=num_simulations, num_workers=50)
        likelihood_estimator = inference.append_simulations(theta, x).train(show_train_summary=True)
        potential_fn, parameter_transform = likelihood_estimator_based_potential(likelihood_estimator, prior, x_o)
        posterior = MCMCPosterior(
            potential_fn=potential_fn, 
            proposal=prior, 
            theta_transform=parameter_transform, 
            num_workers=5, 
            init_strategy="sir", 
            method="slice_np_vectorized",
            num_chains=num_chains)
        proposal = posterior.set_default_x(x_o)
    
    return posterior

posteriors = [get_posterior(x_o=x_o, prior=prior, simulator=linear_gaussian, num_rounds=2, num_simulations=200, num_chains=10) for _ in range(2)]

ensemble = NeuralPosteriorEnsemble(posteriors)
ensemble.set_default_x(x_o)
best = ensemble.map(show_progress_bars=True)
```

I thought it would be possible to get the MAP estimate even when using `NeuralPosteriorEnsemble` and multiple observations as  `sample` work just fine.


Discussion:

Resolving patch:
diff --git a/sbi/utils/conditional_density_utils.py b/sbi/utils/conditional_density_utils.py
--- a/sbi/utils/conditional_density_utils.py
+++ b/sbi/utils/conditional_density_utils.py
@@ -272,7 +272,11 @@ def condition_mog(
 
 class ConditionedPotential:
     def __init__(
-        self, potential_fn: Callable, condition: Tensor, dims_to_sample: List[int]
+        self,
+        potential_fn: Callable,
+        condition: Tensor,
+        dims_to_sample: List[int],
+        allow_iid_x: bool = False,
     ):
         r"""
         Return conditional posterior log-probability or $-\infty$ if outside prior.
@@ -288,6 +292,7 @@ def __init__(
         self.condition = condition
         self.dims_to_sample = dims_to_sample
         self.device = self.potential_fn.device
+        self.allow_iid_x = allow_iid_x
 
     def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
         r"""
@@ -316,7 +321,7 @@ def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
     def set_x(self, x_o: Optional[Tensor]):
         """Check the shape of the observed data and, if valid, set it."""
         if x_o is not None:
-            x_o = process_x(x_o, allow_iid_x=False).to(self.device)
+            x_o = process_x(x_o, allow_iid_x=self.allow_iid_x).to(self.device)
         self.potential_fn.set_x(x_o)
 
     @property
diff --git a/sbi/utils/posterior_ensemble.py b/sbi/utils/posterior_ensemble.py
--- a/sbi/utils/posterior_ensemble.py
+++ b/sbi/utils/posterior_ensemble.py
@@ -414,7 +414,9 @@ def allow_iid_x(self) -> bool:
     def set_x(self, x_o: Optional[Tensor]):
         """Check the shape of the observed data and, if valid, set it."""
         if x_o is not None:
-            x_o = process_x(x_o, allow_iid_x=False).to(self.device)
+            x_o = process_x(x_o, allow_iid_x=self.allow_iid_x).to(  # type: ignore
+                self.device
+            )
         self._x_o = x_o
         for comp_potential in self.potential_fns:
             comp_potential.set_x(x_o)


Test patch:
diff --git a/tests/ensemble_test.py b/tests/ensemble_test.py
--- a/tests/ensemble_test.py
+++ b/tests/ensemble_test.py
@@ -18,20 +18,22 @@
 
 @pytest.mark.slow
 @pytest.mark.parametrize(
-    "inference_method",
-    [
-        "SNLE_A",
-        "SNRE_A",
-        "SNPE_C",
-    ],
+    "inference_method, num_trials",
+    (
+        ("SNPE_C", 1),
+        pytest.param("SNPE_C", 10, marks=pytest.mark.xfail),
+        ("SNLE_A", 1),
+        ("SNLE_A", 10),
+        ("SNLE_A", 1),
+        ("SNRE_A", 10),
+    ),
 )
-def test_c2st_posterior_ensemble_on_linearGaussian(inference_method):
+def test_c2st_posterior_ensemble_on_linearGaussian(inference_method, num_trials):
     """Test whether NeuralPosteriorEnsemble infers well a simple example with available
     ground truth.
 
     """
 
-    num_trials = 1
     num_dim = 2
     x_o = zeros(num_trials, num_dim)
     num_samples = 1000
@@ -65,7 +67,7 @@ def test_c2st_posterior_ensemble_on_linearGaussian(inference_method):
     posterior.set_default_x(x_o)
 
     # test sampling and evaluation.
-    if inference_method == "SNLE_A" or inference_method == "SNRE_A":
+    if inference_method in ["SNLE_A", "SNRE_A"]:
         samples = posterior.sample(
             (num_samples,), num_chains=20, method="slice_np_vectorized"
         )