Repository: sbi-dev/sbi

Issue:
Maximum-a-posterior estimate given a MultipleIndependent prior and MCMC sampler
Hello,

I am encountering an error when trying to estimate the maximum-a-posterior given an MCMCPosterior: slice_np_vectorized, and a MultipleIndependent prior: created via `process_prior` given a list of PyTorch distributions like `torch.distributions.HalfNormal(scale=torch.ones(1)*2)` given one single round

The trace of the error: 
```
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[34], line 1
----> 1 best = posterior.map(show_progress_bars=True, force_update=True)

File ~/.../lib/python3.9/site-packages/sbi/inference/posteriors/mcmc_posterior.py:638, in MCMCPosterior.map(self, x, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars, force_update)
    583 def map(
    584     self,
    585     x: Optional[Tensor] = None,
   (...)
    593     force_update: bool = False,
    594 ) -> Tensor:
    595     r"""Returns the maximum-a-posteriori estimate (MAP).
    596 
    597     The method can be interrupted (Ctrl-C) when the user sees that the
   (...)
    636         The MAP estimate.
    637     """
--> 638     return super().map(
    639         x=x,
    640         num_iter=num_iter,
    641         num_to_optimize=num_to_optimize,
    642         learning_rate=learning_rate,
    643         init_method=init_method,
    644         num_init_samples=num_init_samples,
    645         save_best_every=save_best_every,
    646         show_progress_bars=show_progress_bars,
    647         force_update=force_update,
    648     )

File ~/.../lib/python3.9/site-packages/sbi/inference/posteriors/base_posterior.py:215, in NeuralPosterior.map(self, x, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars, force_update)
    213 if self._map is None or force_update:
    214     self.potential_fn.set_x(self.default_x)
--> 215     self._map = self._calculate_map(
    216         num_iter=num_iter,
    217         num_to_optimize=num_to_optimize,
    218         learning_rate=learning_rate,
    219         init_method=init_method,
    220         num_init_samples=num_init_samples,
    221         save_best_every=save_best_every,
    222         show_progress_bars=show_progress_bars,
    223     )
    224 return self._map

File ~/.../lib/python3.9/site-packages/sbi/inference/posteriors/base_posterior.py:172, in NeuralPosterior._calculate_map(self, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars)
    169 else:
    170     raise ValueError
--> 172 return gradient_ascent(
    173     potential_fn=self.potential_fn,
    174     inits=inits,
    175     theta_transform=self.theta_transform,
    176     num_iter=num_iter,
    177     num_to_optimize=num_to_optimize,
    178     learning_rate=learning_rate,
    179     save_best_every=save_best_every,
    180     show_progress_bars=show_progress_bars,
    181 )[0]

File ~/.../lib/python3.9/site-packages/sbi/utils/sbiutils.py:886, in gradient_ascent(potential_fn, inits, theta_transform, num_iter, num_to_optimize, learning_rate, save_best_every, show_progress_bars, interruption_note)
    879 log_probs_of_optimized = potential_fn(
    880     theta_transform.inv(optimize_inits)
    881 )
    882 best_theta_iter = optimize_inits[  # type: ignore
    883     torch.argmax(log_probs_of_optimized)
    884 ]
    885 best_log_prob_iter = potential_fn(
--> 886     theta_transform.inv(best_theta_iter)
    887 )
    888 if best_log_prob_iter > best_log_prob_overall:
    889     best_theta_overall = best_theta_iter.detach().clone()

File ~/.../lib/python3.9/site-packages/torch/distributions/transforms.py:151, in Transform.__call__(self, x)
    147 """
    148 Computes the transform `x => y`.
    149 """
    150 if self._cache_size == 0:
--> 151     return self._call(x)
    152 x_old, y_old = self._cached_x_y
    153 if x is x_old:

File ~/.../lib/python3.9/site-packages/torch/distributions/transforms.py:436, in IndependentTransform._call(self, x)
    434 if x.dim() < self.domain.event_dim:
    435     raise ValueError("Too few dimensions on input")
--> 436 return self.base_transform(x)

File ~/.../lib/python3.9/site-packages/torch/distributions/transforms.py:151, in Transform.__call__(self, x)
    147 """
    148 Computes the transform `x => y`.
    149 """
    150 if self._cache_size == 0:
--> 151     return self._call(x)
    152 x_old, y_old = self._cached_x_y
    153 if x is x_old:

File ~/.../lib/python3.9/site-packages/torch/distributions/transforms.py:1017, in CatTransform._call(self, x)
   1016 def _call(self, x):
-> 1017     assert -x.dim() <= self.dim < x.dim()
   1018     assert x.size(self.dim) == self.length
   1019     yslices = []

AssertionError: 
```

I think that this error is related to #650 (this time with MCMC). However, I'm unable to grasp the details discussed, making it difficult for me to detect the problem in my code.


Discussion:
Thanks for reporting, this indeed sounds like a bug. As a quick workaround, you could turn off the transformation for the `.map()` computation:
```python
from sbi.inference import likelihood_estimator_based_potential, MCMCPosterior

likelihood_estimator = inference.append_simulations(theta, x).train()

potential_fn, parameter_transform = likelihood_estimator_based_potential(
    likelihood_estimator, prior, x_o
)
posterior = MCMCPosterior(
    potential_fn, proposal=prior, theta_transform=None  # instead of `parameter_transform`
)
map_estimate = posterior.set_default_x(x_o).map()
```
Thank you for your response, as you suggested I set `theta_transform` to `None`, and got a new posterior but when calling `.map()` I get the following error:

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[12], line 1
----> 1 map_estimate = posterior.map(show_progress_bars=True)

File ~/linking_game/LinkingGameModel/linking_game_env/lib/python3.9/site-packages/sbi/inference/posteriors/mcmc_posterior.py:638, in MCMCPosterior.map(self, x, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars, force_update)
    583 def map(
    584     self,
    585     x: Optional[Tensor] = None,
   (...)
    593     force_update: bool = False,
    594 ) -> Tensor:
    595     r"""Returns the maximum-a-posteriori estimate (MAP).
    596 
    597     The method can be interrupted (Ctrl-C) when the user sees that the
   (...)
    636         The MAP estimate.
    637     """
--> 638     return super().map(
    639         x=x,
    640         num_iter=num_iter,
    641         num_to_optimize=num_to_optimize,
    642         learning_rate=learning_rate,
    643         init_method=init_method,
    644         num_init_samples=num_init_samples,
    645         save_best_every=save_best_every,
    646         show_progress_bars=show_progress_bars,
    647         force_update=force_update,
    648     )

File ~/linking_game/LinkingGameModel/linking_game_env/lib/python3.9/site-packages/sbi/inference/posteriors/base_posterior.py:215, in NeuralPosterior.map(self, x, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars, force_update)
    213 if self._map is None or force_update:
    214     self.potential_fn.set_x(self.default_x)
--> 215     self._map = self._calculate_map(
    216         num_iter=num_iter,
    217         num_to_optimize=num_to_optimize,
    218         learning_rate=learning_rate,
    219         init_method=init_method,
    220         num_init_samples=num_init_samples,
    221         save_best_every=save_best_every,
    222         show_progress_bars=show_progress_bars,
    223     )
    224 return self._map

File ~/linking_game/LinkingGameModel/linking_game_env/lib/python3.9/site-packages/sbi/inference/posteriors/base_posterior.py:172, in NeuralPosterior._calculate_map(self, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars)
    169 else:
    170     raise ValueError
--> 172 return gradient_ascent(
    173     potential_fn=self.potential_fn,
    174     inits=inits,
    175     theta_transform=self.theta_transform,
    176     num_iter=num_iter,
    177     num_to_optimize=num_to_optimize,
    178     learning_rate=learning_rate,
    179     save_best_every=save_best_every,
    180     show_progress_bars=show_progress_bars,
    181 )[0]

File ~/.../lib/python3.9/site-packages/sbi/utils/sbiutils.py:879, in gradient_ascent(potential_fn, inits, theta_transform, num_iter, num_to_optimize, learning_rate, save_best_every, show_progress_bars, interruption_note)
    876 with torch.no_grad():
    877     if iter_ % save_best_every == 0 or iter_ == num_iter - 1:
    878         # Evaluate the optimized locations and pick the best one.
--> 879         log_probs_of_optimized = potential_fn(
    880             theta_transform.inv(optimize_inits)
    881         )
    882         best_theta_iter = optimize_inits[  # type: ignore
    883             torch.argmax(log_probs_of_optimized)
    884         ]
    885         best_log_prob_iter = potential_fn(
    886             theta_transform.inv(best_theta_iter)
    887         )

File ~/.../lib/python3.9/site-packages/sbi/inference/potentials/likelihood_based_potential.py:99, in LikelihoodBasedPotential.__call__(self, theta, track_gradients)
     91 # Calculate likelihood over trials and in one batch.
     92 log_likelihood_trial_sum = _log_likelihoods_over_trials(
     93     x=self.x_o,
     94     theta=theta.to(self.device),
     95     net=self.likelihood_estimator,
     96     track_gradients=track_gradients,
     97 )
---> 99 return log_likelihood_trial_sum + self.prior.log_prob(theta)

File ~/.../lib/python3.9/site-packages/sbi/utils/user_input_checks_utils.py:314, in MultipleIndependent.log_prob(self, value)
    312     v = value[:, dims_covered : dims_covered + ndims]
    313     # Reshape here to ensure all returned log_probs are 2D for concatenation.
--> 314     log_probs.append(d.log_prob(v).reshape(num_samples, 1))
    315     dims_covered += ndims
    317 # Sum accross last dimension to get joint log prob over all distributions.

File ~/.../lib/python3.9/site-packages/pyro/distributions/torch_patch.py:65, in _HalfCauchy_logprob(self, value)
     62 @patch_dependency("torch.distributions.HalfCauchy.log_prob")
     63 def _HalfCauchy_logprob(self, value):
     64     if self._validate_args:
---> 65         self._validate_sample(value)
     66     value = torch.as_tensor(
     67         value, dtype=self.base_dist.scale.dtype, device=self.base_dist.scale.device
     68     )
     69     log_prob = self.base_dist.log_prob(value) + math.log(2)

File ~/.../lib/python3.9/site-packages/torch/distributions/distribution.py:294, in Distribution._validate_sample(self, value)
    292 valid = support.check(value)
    293 if not valid.all():
--> 294     raise ValueError(
    295         "Expected value argument "
    296         f"({type(value).__name__} of shape {tuple(value.shape)}) "
    297         f"to be within the support ({repr(support)}) "
    298         f"of the distribution {repr(self)}, "
    299         f"but found invalid values:\n{value}"
    300     )

ValueError: Expected value argument (Tensor of shape (100, 1)) to be within the support (GreaterThanEq(lower_bound=0.0)) of the distribution HalfCauchy(), but found invalid values:
tensor([[ 4.9898e+00],
	[ 4.7333e-01],
	...
	])
```
Ah damn. Can you try setting `.map(num_init_samples=1)` (or some other small value)?
Yes, it did work (consistently with up to `num_init_samples=5`) but doesn't this mean that the gradient descent will start with a very low number of initial points and that some calls to `map` will fail when the sample prawn from the posterior is outside the support of the prior distribution?
The map is found by performing gradient ascent on `posterior.potential()`. So yes, by using a low number of initial values you risk that all of them run into local maxima but do not actually find the global MAP.

You could also try `.map(num_init_samples=10_000, num_to_optimize=1)`. This initially samples 10k times from the prior, sorts them according `posterior.potential()` and then runs gradient ascent only for the best one.

Anyways, we will have to fix this (it's def a bug) but I won't get to it today.
In doing so it happens that it occasionally selects values beyond the distribution's support:

```
Optimizing MAP estimate. Iterations: 9 / 1000. Performance in iteration 0: -1200.72 (= unnormalized log-prob)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[24], line 1
----> 1 best = posterior.map(num_init_samples=10_000, num_to_optimize=1, show_progress_bars=True, force_update=True)

File ~/.../lib/python3.9/site-packages/sbi/inference/posteriors/mcmc_posterior.py:638, in MCMCPosterior.map(self, x, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars, force_update)
    583 def map(
    584     self,
    585     x: Optional[Tensor] = None,
   (...)
    593     force_update: bool = False,
    594 ) -> Tensor:
    595     r"""Returns the maximum-a-posteriori estimate (MAP).
    596 
    597     The method can be interrupted (Ctrl-C) when the user sees that the
   (...)
    636         The MAP estimate.
    637     """
--> 638     return super().map(
    639         x=x,
    640         num_iter=num_iter,
    641         num_to_optimize=num_to_optimize,
    642         learning_rate=learning_rate,
    643         init_method=init_method,
    644         num_init_samples=num_init_samples,
    645         save_best_every=save_best_every,
    646         show_progress_bars=show_progress_bars,
    647         force_update=force_update,
    648     )

File ~/.../lib/python3.9/site-packages/sbi/inference/posteriors/base_posterior.py:215, in NeuralPosterior.map(self, x, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars, force_update)
    213 if self._map is None or force_update:
    214     self.potential_fn.set_x(self.default_x)
--> 215     self._map = self._calculate_map(
    216         num_iter=num_iter,
    217         num_to_optimize=num_to_optimize,
    218         learning_rate=learning_rate,
    219         init_method=init_method,
    220         num_init_samples=num_init_samples,
    221         save_best_every=save_best_every,
    222         show_progress_bars=show_progress_bars,
    223     )
    224 return self._map

File ~/.../lib/python3.9/site-packages/sbi/inference/posteriors/base_posterior.py:172, in NeuralPosterior._calculate_map(self, num_iter, num_to_optimize, learning_rate, init_method, num_init_samples, save_best_every, show_progress_bars)
    169 else:
    170     raise ValueError
--> 172 return gradient_ascent(
    173     potential_fn=self.potential_fn,
    174     inits=inits,
    175     theta_transform=self.theta_transform,
    176     num_iter=num_iter,
    177     num_to_optimize=num_to_optimize,
    178     learning_rate=learning_rate,
    179     save_best_every=save_best_every,
    180     show_progress_bars=show_progress_bars,
    181 )[0]

File ~/.../lib/python3.9/site-packages/sbi/utils/sbiutils.py:871, in gradient_ascent(potential_fn, inits, theta_transform, num_iter, num_to_optimize, learning_rate, save_best_every, show_progress_bars, interruption_note)
    869 while iter_ < num_iter:
    870     optimizer.zero_grad()
--> 871     probs = potential_fn(theta_transform.inv(optimize_inits)).squeeze()
    872     loss = -probs.sum()
    873     loss.backward()

File ~/.../lib/python3.9/site-packages/sbi/inference/potentials/likelihood_based_potential.py:99, in LikelihoodBasedPotential.__call__(self, theta, track_gradients)
     91 # Calculate likelihood over trials and in one batch.
     92 log_likelihood_trial_sum = _log_likelihoods_over_trials(
     93     x=self.x_o,
     94     theta=theta.to(self.device),
     95     net=self.likelihood_estimator,
     96     track_gradients=track_gradients,
     97 )
---> 99 return log_likelihood_trial_sum + self.prior.log_prob(theta)

File ~/.../lib/python3.9/site-packages/sbi/utils/user_input_checks_utils.py:314, in MultipleIndependent.log_prob(self, value)
    312     v = value[:, dims_covered : dims_covered + ndims]
    313     # Reshape here to ensure all returned log_probs are 2D for concatenation.
--> 314     log_probs.append(d.log_prob(v).reshape(num_samples, 1))
    315     dims_covered += ndims
    317 # Sum accross last dimension to get joint log prob over all distributions.

File ~/.../lib/python3.9/site-packages/pyro/distributions/torch_patch.py:65, in _HalfCauchy_logprob(self, value)
     62 @patch_dependency("torch.distributions.HalfCauchy.log_prob")
     63 def _HalfCauchy_logprob(self, value):
     64     if self._validate_args:
---> 65         self._validate_sample(value)
     66     value = torch.as_tensor(
     67         value, dtype=self.base_dist.scale.dtype, device=self.base_dist.scale.device
     68     )
     69     log_prob = self.base_dist.log_prob(value) + math.log(2)

File ~/.../lib/python3.9/site-packages/torch/distributions/distribution.py:294, in Distribution._validate_sample(self, value)
    292 valid = support.check(value)
    293 if not valid.all():
--> 294     raise ValueError(
    295         "Expected value argument "
    296         f"({type(value).__name__} of shape {tuple(value.shape)}) "
    297         f"to be within the support ({repr(support)}) "
    298         f"of the distribution {repr(self)}, "
    299         f"but found invalid values:\n{value}"
    300     )

ValueError: Expected value argument (Tensor of shape (1, 1)) to be within the support (GreaterThanEq(lower_bound=0.0)) of the distribution HalfCauchy(), but found invalid values:
tensor([[-0.0051]], grad_fn=)
```

Resolving patch:
diff --git a/sbi/utils/sbiutils.py b/sbi/utils/sbiutils.py
--- a/sbi/utils/sbiutils.py
+++ b/sbi/utils/sbiutils.py
@@ -349,7 +349,10 @@ def warn_on_iid_x(num_trials):
 
 
 def check_warn_and_setstate(
-    state_dict: Dict, key_name: str, replacement_value: Any, warning_msg: str = ""
+    state_dict: Dict,
+    key_name: str,
+    replacement_value: Any,
+    warning_msg: str = "",
 ) -> Tuple[Dict, str]:
     """
     Check if `key_name` is in `state_dict` and add it if not.
@@ -881,7 +884,7 @@ def gradient_ascent(
                     )
                     best_theta_iter = optimize_inits[  # type: ignore
                         torch.argmax(log_probs_of_optimized)
-                    ]
+                    ].view(1, -1)
                     best_log_prob_iter = potential_fn(
                         theta_transform.inv(best_theta_iter)
                     )
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -17,9 +17,7 @@
 # Package meta-data.
 NAME = "sbi"
 DESCRIPTION = "Simulation-based inference."
-KEYWORDS = (
-    "bayesian parameter inference system_identification simulator PyTorch"
-)
+KEYWORDS = "bayesian parameter inference system_identification simulator PyTorch"
 URL = "https://github.com/mackelab/sbi"
 EMAIL = "sbi@mackelab.org"
 AUTHOR = """Álvaro Tejero-Cantero, Jakob H. Macke, Jan-Matthis Lückmann, Conor M.
@@ -106,9 +104,7 @@ def run(self):
             pass
 
         self.status("Building Source and Wheel (universal) distribution…")
-        os.system(
-            "{0} setup.py sdist bdist_wheel --universal".format(sys.executable)
-        )
+        os.system("{0} setup.py sdist bdist_wheel --universal".format(sys.executable))
 
         self.status("Uploading the package to PyPI via Twine…")
         os.system("twine upload dist/*")
@@ -131,9 +127,7 @@ def run(self):
     author_email=EMAIL,
     python_requires=REQUIRES_PYTHON,
     url=URL,
-    packages=find_packages(
-        exclude=["tests", "*.tests", "*.tests.*", "tests.*"]
-    ),
+    packages=find_packages(exclude=["tests", "*.tests", "*.tests.*", "tests.*"]),
     install_requires=REQUIRED,
     extras_require=EXTRAS,
     include_package_data=True,


Test patch:
diff --git a/tests/linearGaussian_snle_test.py b/tests/linearGaussian_snle_test.py
--- a/tests/linearGaussian_snle_test.py
+++ b/tests/linearGaussian_snle_test.py
@@ -4,10 +4,10 @@
 from __future__ import annotations
 
 import pytest
+import torch
 from torch import eye, ones, zeros
-from torch.distributions import MultivariateNormal
+from torch.distributions import HalfNormal, MultivariateNormal
 
-from sbi import utils as utils
 from sbi.inference import (
     SNLE,
     ImportanceSamplingPosterior,
@@ -25,7 +25,7 @@
     samples_true_posterior_linear_gaussian_uniform_prior,
     true_posterior_linear_gaussian_mvn_prior,
 )
-from sbi.utils import likelihood_nn
+from sbi.utils import BoxUniform, likelihood_nn, process_prior
 
 from .test_utils import check_c2st, get_prob_outside_uniform_prior
 
@@ -103,7 +103,10 @@ def test_c2st_snl_on_linearGaussian(density_estimator="maf"):
     )
     simulator, prior = prepare_for_sbi(
         lambda theta: linear_gaussian(
-            theta, likelihood_shift, likelihood_cov, num_discarded_dims=discard_dims
+            theta,
+            likelihood_shift,
+            likelihood_cov,
+            num_discarded_dims=discard_dims,
         ),
         prior,
     )
@@ -156,10 +159,11 @@ def test_c2st_and_map_snl_on_linearGaussian_different(num_dim: int, prior_str: s
         prior_cov = eye(num_dim)
         prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
     else:
-        prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
+        prior = BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
     density_estimator = likelihood_nn("maf", num_transforms=3)
     inference = SNLE(density_estimator=density_estimator, show_progress_bars=False)
@@ -204,11 +208,15 @@ def test_c2st_and_map_snl_on_linearGaussian_different(num_dim: int, prior_str: s
 
         # Check performance based on c2st accuracy.
         check_c2st(
-            samples, target_samples, alg=f"snle_a-{prior_str}-prior-{num_trials}-trials"
+            samples,
+            target_samples,
+            alg=f"snle_a-{prior_str}-prior-{num_trials}-trials",
         )
 
         map_ = posterior.map(
-            num_init_samples=1_000, init_method="proposal", show_progress_bars=False
+            num_init_samples=1_000,
+            init_method="proposal",
+            show_progress_bars=False,
         )
 
         # TODO: we do not have a test for SNL log_prob(). This is because the output
@@ -225,6 +233,43 @@ def test_c2st_and_map_snl_on_linearGaussian_different(num_dim: int, prior_str: s
             assert ((map_ - gt_posterior.mean) ** 2).sum() < 0.5
 
 
+@pytest.mark.parametrize("use_transform", (True, False))
+def test_map_with_multiple_independent_prior(use_transform):
+    """Test whether map works with multiple independent priors, see issue #841, #650."""
+
+    dim = 2
+    prior, *_ = process_prior(
+        [
+            BoxUniform(low=-ones(dim), high=ones(dim)),
+            HalfNormal(scale=ones(1) * 2),
+        ]
+    )
+
+    def simulator(theta):
+        return theta[:, 2:] * torch.randn_like(theta[:, :2]) + theta[:, :2]
+
+    num_simulations = 1000
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
+    x_o = zeros((1, dim))
+
+    trainer = SNLE(prior).append_simulations(theta, x)
+    likelihood_estimator = trainer.train(max_num_epochs=5)
+
+    potential_fn, parameter_transform = likelihood_estimator_based_potential(
+        likelihood_estimator,
+        prior,
+        x_o=x_o,
+    )
+    posterior = MCMCPosterior(
+        potential_fn,
+        proposal=prior,
+        theta_transform=parameter_transform if use_transform else None,
+    )
+    posterior.map()
+    posterior.set_default_x(x_o).map(num_iter=10)
+
+
 @pytest.mark.slow
 @pytest.mark.parametrize("num_trials", (1, 3))
 def test_c2st_multi_round_snl_on_linearGaussian(num_trials: int):
@@ -248,7 +293,8 @@ def test_c2st_multi_round_snl_on_linearGaussian(num_trials: int):
     target_samples = gt_posterior.sample((num_samples,))
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
     inference = SNLE(show_progress_bars=False)
 
@@ -268,7 +314,10 @@ def test_c2st_multi_round_snl_on_linearGaussian(num_trials: int):
     )
 
     theta, x = simulate_for_sbi(
-        simulator, posterior1, num_simulations_per_round, simulation_batch_size=50
+        simulator,
+        posterior1,
+        num_simulations_per_round,
+        simulation_batch_size=50,
     )
     likelihood_estimator = inference.append_simulations(theta, x).train()
     potential_fn, theta_transform = likelihood_estimator_based_potential(
@@ -311,7 +360,8 @@ def test_c2st_multi_round_snl_on_linearGaussian_vi(num_trials: int):
     target_samples = gt_posterior.sample((num_samples,))
 
     simulator, prior = prepare_for_sbi(
-        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov),
+        prior,
     )
     inference = SNLE(show_progress_bars=False)
 
@@ -329,7 +379,10 @@ def test_c2st_multi_round_snl_on_linearGaussian_vi(num_trials: int):
     posterior1.train()
 
     theta, x = simulate_for_sbi(
-        simulator, posterior1, num_simulations_per_round, simulation_batch_size=50
+        simulator,
+        posterior1,
+        num_simulations_per_round,
+        simulation_batch_size=50,
     )
     likelihood_estimator = inference.append_simulations(theta, x).train()
     potential_fn, theta_transform = likelihood_estimator_based_potential(
@@ -408,7 +461,7 @@ def test_api_snl_sampling_methods(
     if prior_str == "gaussian":
         prior = MultivariateNormal(loc=zeros(num_dim), covariance_matrix=eye(num_dim))
     else:
-        prior = utils.BoxUniform(-1.0 * ones(num_dim), ones(num_dim))
+        prior = BoxUniform(-1.0 * ones(num_dim), ones(num_dim))
 
     # Why do we have this if-case? Only the `MCMCPosterior` uses the `init_strategy`.
     # Thus, we would not like to run, e.g., VI with all init_strategies, but only once
@@ -450,7 +503,9 @@ def test_api_snl_sampling_methods(
             )
         else:
             posterior = VIPosterior(
-                potential_fn, theta_transform=theta_transform, vi_method=sampling_method
+                potential_fn,
+                theta_transform=theta_transform,
+                vi_method=sampling_method,
             )
             posterior.train(max_num_iters=10)