Repository: sbi-dev/sbi

Issue:
Density Estimator batched sample mixes up samples from different posteriors
**Describe the bug**
Given a batched observation, i.e., x1 and x2, the sampling method mixes up samples from different distributions.

**To Reproduce**
```python
import torch

from sbi import analysis as analysis
from sbi import utils as utils
from sbi.inference.base import infer

num_dim = 3
prior = utils.BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))

def simulator(parameter_set):
    return 1.0 + parameter_set + torch.randn(parameter_set.shape) * 0.1

posterior = infer(simulator, prior, method="SNPE", num_simulations=200)
observation = torch.stack([torch.zeros(3), torch.ones(3)])
posterior_samples = posterior.posterior_estimator.sample((1000,), condition=observation)

# Outputs an multimodal distribution, but should be unimodal (mixes up samples from the two different x_os)
samples1 = posterior_samples[:,0].detach()
_ = analysis.pairplot([samples1], limits=[[-2, 2], [-2, 2], [-2, 2]], figsize=(6, 6))
```

**Additional context**

Likely a "reshaping" bug. 



Discussion:

Resolving patch:
diff --git a/sbi/neural_nets/density_estimators/nflows_flow.py b/sbi/neural_nets/density_estimators/nflows_flow.py
--- a/sbi/neural_nets/density_estimators/nflows_flow.py
+++ b/sbi/neural_nets/density_estimators/nflows_flow.py
@@ -135,12 +135,8 @@ def sample(self, sample_shape: Shape, condition: Tensor) -> Tensor:
         num_samples = torch.Size(sample_shape).numel()
 
         samples = self.net.sample(num_samples, context=condition)
-
-        return samples.reshape((
-            *sample_shape,
-            condition_batch_dim,
-            -1,
-        ))
+        samples = samples.transpose(0, 1)
+        return samples.reshape((*sample_shape, condition_batch_dim, *self.input_shape))
 
     def sample_and_log_prob(
         self, sample_shape: torch.Size, condition: Tensor, **kwargs


Test patch:
diff --git a/tests/density_estimator_test.py b/tests/density_estimator_test.py
--- a/tests/density_estimator_test.py
+++ b/tests/density_estimator_test.py
@@ -283,6 +283,75 @@ def test_correctness_of_density_estimator_log_prob(
     assert torch.allclose(log_probs[0, :], log_probs[1, :])
 
 
+@pytest.mark.parametrize(
+    "density_estimator_build_fn",
+    (
+        build_mdn,
+        build_maf,
+        build_maf_rqs,
+        build_nsf,
+        build_zuko_bpf,
+        build_zuko_gf,
+        build_zuko_maf,
+        build_zuko_naf,
+        build_zuko_ncsf,
+        build_zuko_nice,
+        build_zuko_nsf,
+        build_zuko_sospf,
+        build_zuko_unaf,
+        build_categoricalmassestimator,
+        build_mnle,
+    ),
+)
+@pytest.mark.parametrize("input_event_shape", ((1,), (4,)))
+@pytest.mark.parametrize("condition_event_shape", ((1,), (7,)))
+def test_correctness_of_batched_vs_seperate_sample_and_log_prob(
+    density_estimator_build_fn, input_event_shape, condition_event_shape
+):
+    input_sample_dim = 2
+    batch_dim = 2
+    density_estimator, inputs, condition = _build_density_estimator_and_tensors(
+        density_estimator_build_fn,
+        input_event_shape,
+        condition_event_shape,
+        batch_dim,
+        input_sample_dim,
+    )
+    # Batched vs separate sampling
+    samples = density_estimator.sample((1000,), condition=condition)
+    samples_separate1 = density_estimator.sample(
+        (1000,), condition=condition[0][None, ...]
+    )
+    samples_separate2 = density_estimator.sample(
+        (1000,), condition=condition[1][None, ...]
+    )
+
+    # Check if means are approx. same
+    samples_m = torch.mean(samples, dim=0, dtype=torch.float32)
+    samples_separate1_m = torch.mean(samples_separate1, dim=0, dtype=torch.float32)
+    samples_separate2_m = torch.mean(samples_separate2, dim=0, dtype=torch.float32)
+    samples_sep_m = torch.cat([samples_separate1_m, samples_separate2_m], dim=0)
+
+    assert torch.allclose(
+        samples_m, samples_sep_m, atol=0.5, rtol=0.5
+    ), "Batched sampling is not consistent with separate sampling."
+
+    # Batched vs separate log_prob
+    log_probs = density_estimator.log_prob(inputs, condition=condition)
+
+    log_probs_separate1 = density_estimator.log_prob(
+        inputs[:, :1], condition=condition[0][None, ...]
+    )
+    log_probs_separate2 = density_estimator.log_prob(
+        inputs[:, 1:], condition=condition[1][None, ...]
+    )
+    log_probs_sep = torch.hstack([log_probs_separate1, log_probs_separate2])
+
+    assert torch.allclose(
+        log_probs, log_probs_sep, atol=1e-2, rtol=1e-2
+    ), "Batched log_prob is not consistent with separate log_prob."
+
+
 def _build_density_estimator_and_tensors(
     density_estimator_build_fn: str,
     input_event_shape: Tuple[int],