Repository: sbi-dev/sbi

Issue:
MCMC breaks with multi dimensional x
If `x` is multidimensional, e.g., images with shape `32x32`, this will break the current version of `slice_np` and `slice_np_vectorized` and it will result in endless loops in the `pyro` mcmc versions. 

The `numpy` slice problem is fixed quickly, it is just about the repeating here:
https://github.com/mackelab/sbi/blob/6b5ed7be1d7522546b06c39aec1f206a354cc2ef/sbi/inference/posteriors/direct_posterior.py#L468

It has to be extended for every dimension of `x`. I will make a PR.

As for the `pyro` problem, I don't know.


Discussion:

Resolving patch:
diff --git a/sbi/inference/posteriors/direct_posterior.py b/sbi/inference/posteriors/direct_posterior.py
--- a/sbi/inference/posteriors/direct_posterior.py
+++ b/sbi/inference/posteriors/direct_posterior.py
@@ -465,10 +465,17 @@ def np_potential(self, theta: np.ndarray) -> ScalarFloat:
         theta = torch.as_tensor(theta, dtype=torch.float32)
         theta = ensure_theta_batched(theta)
         num_batch = theta.shape[0]
-        x = ensure_x_batched(self.x).repeat(num_batch, 1)
+
+        x_batched = ensure_x_batched(self.x)
+        # Repeat x over batch dim to match theta batch, accounting for multi-D x.
+        x_repeated = x_batched.repeat(
+            num_batch, *(1 for _ in range(x_batched.ndim - 1))
+        )
 
         with torch.set_grad_enabled(False):
-            target_log_prob = self.posterior_nn.log_prob(inputs=theta, context=self.x,)
+            target_log_prob = self.posterior_nn.log_prob(
+                inputs=theta, context=x_repeated,
+            )
             is_within_prior = torch.isfinite(self.prior.log_prob(theta))
             target_log_prob[~is_within_prior] = -float("Inf")
 
diff --git a/sbi/inference/posteriors/ratio_based_posterior.py b/sbi/inference/posteriors/ratio_based_posterior.py
--- a/sbi/inference/posteriors/ratio_based_posterior.py
+++ b/sbi/inference/posteriors/ratio_based_posterior.py
@@ -327,10 +327,20 @@ def np_potential(self, theta: np.array) -> ScalarFloat:
         theta = torch.as_tensor(theta, dtype=torch.float32)
         theta = ensure_theta_batched(theta)
         num_batch = theta.shape[0]
-        x = ensure_x_batched(self.x).repeat(num_batch, 1)
+        x_batched = ensure_x_batched(self.x)
+        # Repeat x over batch dim to match theta batch, accounting for multi-D x.
+        x_repeated = x_batched.repeat(
+            num_batch, *(1 for _ in range(x_batched.ndim - 1))
+        )
 
+        assert (
+            x_batched.ndim == 2
+        ), """X must not be multidimensional for ratio-based methods because it will be
+              concatenated with theta."""
         with torch.set_grad_enabled(False):
-            log_ratio = self.classifier(torch.cat((theta, x), dim=1)).reshape(-1)
+            log_ratio = self.classifier(torch.cat((theta, x_repeated), dim=1)).reshape(
+                -1
+            )
 
         # Notice opposite sign to pyro potential.
         return log_ratio + self.prior.log_prob(theta)


Test patch:
diff --git a/tests/multidimensional_x_test.py b/tests/multidimensional_x_test.py
--- a/tests/multidimensional_x_test.py
+++ b/tests/multidimensional_x_test.py
@@ -71,13 +71,34 @@ def test_inference_with_2d_x(embedding, method):
     x_o = simulator(theta_o)
 
     if method == SNPE:
-        net_provider = utils.posterior_nn(model="mdn", embedding_net=embedding(),)
+        kwargs = dict(
+            density_estimator=utils.posterior_nn(
+                model="mdn", embedding_net=embedding(),
+            ),
+            sample_with_mcmc=True,
+        )
     elif method == SNLE:
-        net_provider = utils.likelihood_nn(model="mdn", embedding_net=embedding())
+        kwargs = dict(
+            density_estimator=utils.likelihood_nn(
+                model="mdn", embedding_net=embedding()
+            )
+        )
     else:
-        net_provider = utils.classifier_nn(model="mlp", embedding_net_x=embedding(),)
-
-    infer = method(simulator, prior, 1, 1, net_provider, show_progress_bars=False,)
+        kwargs = dict(
+            density_estimator=utils.classifier_nn(
+                model="mlp", embedding_net_x=embedding(),
+            )
+        )
+
+    infer = method(
+        simulator,
+        prior,
+        1,
+        1,
+        show_progress_bars=False,
+        mcmc_method="slice_np",
+        **kwargs,
+    )
 
     posterior = infer(
         num_simulations=num_simulations, training_batch_size=100, max_num_epochs=10