Repository: sbi-dev/sbi

Issue:
SNLE + multiple independent prior + MCMC with num_workers > 1 throws error
SNLE:
- multiple independent prior + num_workers=2 leads to an error (ValueError: Expected value argument to be within the support of the distribution Uniform)
- no error if either num_workers=1 or using a BoxUniform prior

SNPE:
- seems to work

```
from sbi.inference.snpe import SNPE_A
from sbi.inference.snle import SNLE_A
from sbi.utils import BoxUniform
from sbi.inference.posteriors import MCMCPosterior
from sbi.inference.potentials import posterior_estimator_based_potential, likelihood_estimator_based_potential
from sbi.utils.user_input_checks import process_prior
from sbi.analysis.sbc import run_sbc
import torch

default_params = torch.rand(size=(4, 3),dtype=torch.float32)
default_obs = torch.rand(size=(4, 1), dtype=torch.float32)

# multiple independent prior
prior = [torch.distributions.Uniform(torch.FloatTensor([0]), torch.FloatTensor([20])),
        torch.distributions.Uniform(torch.FloatTensor([-10]), torch.FloatTensor([10])),
        torch.distributions.Uniform(torch.FloatTensor([0.5]), torch.FloatTensor([3]))]
# box uniform prior
# prior = BoxUniform(torch.Tensor([0, -10, 0.5]), torch.Tensor([20, 10, 3]))
prior, _, _ = process_prior(prior)

# inference = SNPE_A()
# density_estimator = inference.append_simulations(default_params, default_obs).train()
# potential_fn, theta_transform = posterior_estimator_based_potential(density_estimator, prior, default_obs[0])

inference = SNLE_A()
density_estimator = inference.append_simulations(default_params, default_obs).train()
potential_fn, theta_transform = likelihood_estimator_based_potential(density_estimator, prior, default_obs[0])

posterior = MCMCPosterior(potential_fn, proposal=prior, theta_transform=theta_transform)

# this line throws an error
ranks, dap_samples = run_sbc(default_params, default_obs, posterior, num_posterior_samples=10, num_workers=2, sbc_batch_size=2)
```



Discussion:

Resolving patch:
diff --git a/sbi/samplers/mcmc/init_strategy.py b/sbi/samplers/mcmc/init_strategy.py
--- a/sbi/samplers/mcmc/init_strategy.py
+++ b/sbi/samplers/mcmc/init_strategy.py
@@ -63,9 +63,8 @@ def sir(
         init_param_candidates = []
         for i in range(sir_num_batches):
             batch_draws = proposal.sample((sir_batch_size,)).detach()
-            transformed_batch_draws = transform(batch_draws)
-            init_param_candidates.append(transformed_batch_draws)
-            log_weights.append(potential_fn(transformed_batch_draws).detach())
+            init_param_candidates.append(batch_draws)
+            log_weights.append(potential_fn(batch_draws).detach())
         log_weights = torch.cat(log_weights)
         init_param_candidates = torch.cat(init_param_candidates)
 
@@ -77,4 +76,5 @@ def sir(
         probs /= probs.sum()
 
         idxs = torch.multinomial(probs, 1, replacement=False)
-        return init_param_candidates[idxs, :]
+        # Return transformed sample.
+        return transform(init_param_candidates[idxs, :])


Test patch:
diff --git a/tests/sbc_test.py b/tests/sbc_test.py
--- a/tests/sbc_test.py
+++ b/tests/sbc_test.py
@@ -6,29 +6,34 @@
 import pytest
 import torch
 from torch import eye, ones, zeros
-from torch.distributions import MultivariateNormal
+from torch.distributions import MultivariateNormal, Uniform
 
 from sbi.analysis import check_sbc, get_nltp, run_sbc
-from sbi.inference import SNPE_C, simulate_for_sbi
+from sbi.inference import SNLE, SNPE, simulate_for_sbi
 from sbi.simulators import linear_gaussian
+from sbi.utils import BoxUniform, MultipleIndependent
 
 
-def test_running_sbc(method=SNPE_C, model="mdn"):
+@pytest.mark.parametrize("prior", ("boxuniform", "independent"))
+@pytest.mark.parametrize("method", (SNPE, SNLE))
+def test_running_sbc(method, prior, model="mdn"):
     """Tests running inference and then SBC and obtaining nltp."""
 
     num_dim = 2
+    if prior == "boxuniform":
+        prior = BoxUniform(-torch.ones(num_dim), torch.ones(num_dim))
+    else:
+        prior = MultipleIndependent(
+            [Uniform(-torch.ones(1), torch.ones(1)) for _ in range(num_dim)]
+        )
+
     num_simulations = 100
     max_num_epochs = 1
     num_sbc_runs = 2
 
-    x_o = zeros(1, num_dim)
     likelihood_shift = -1.0 * ones(num_dim)
     likelihood_cov = 0.3 * eye(num_dim)
 
-    prior_mean = zeros(num_dim)
-    prior_cov = eye(num_dim)
-    prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
-
     def simulator(theta):
         return linear_gaussian(theta, likelihood_shift, likelihood_cov)
 
@@ -39,12 +44,12 @@ def simulator(theta):
     _ = inferer.append_simulations(theta, x).train(
         training_batch_size=100, max_num_epochs=max_num_epochs
     )
-    posterior = inferer.build_posterior().set_default_x(x_o)
+    posterior = inferer.build_posterior()
 
     thetas = prior.sample((num_sbc_runs,))
     xs = simulator(thetas)
 
-    run_sbc(thetas, xs, posterior, num_workers=5)
+    run_sbc(thetas, xs, posterior, num_workers=1, num_posterior_samples=10)
 
     # Check nltp
     get_nltp(thetas, xs, posterior)