Repository: sbi-dev/sbi

Issue:
Sensitivity on a posterior trained on GPU
Hi,

I am trying the sensitivity analysis feature and when I try to compute the eigenvalues and vectors using a posterior trained on the GPU, I get a RuntimeError that says: Expected all tensors to be on the same device, but found at least two devices.

The callback is: sensitivity_analysis.py (line 430), then a sequence of torch files (container.py and module.py), and finally sbiutils.py (line 115)

Could it be because the ActiveSubspace object is on cpu and the posterior on GPU?

Edit: I found a way to work around the issue by simply adding .to('cpu') on the posterior samples (line 423 of sensitivity_analysis.py ) the error from line 430 disappeared, which led to a new error on line 444 which I worked around by also adding .to('cpu') on the prior_scale of that same line (444). All it does is ensure that the involved tensors are on the cpu for the "find_direction" function. so it's not a fix, but it does confirm that the issue comes from the posterior and the ActiveSubspace being on different devices.




Discussion:
Hi there,

yeah, that sounds very much like a bug, I had never used the `ActiveSubspace` with GPU. Great that you found a workaround! I will leave this issue open until I fix the GPU issue.

Thanks for reporting!

Michael
I've just come across the same error in the `conditional_pairplot` utility too. I can provide a full stack trace if that's useful

Resolving patch:
diff --git a/sbi/analysis/plot.py b/sbi/analysis/plot.py
--- a/sbi/analysis/plot.py
+++ b/sbi/analysis/plot.py
@@ -639,6 +639,7 @@ def conditional_pairplot(
 
     Returns: figure and axis of posterior distribution plot
     """
+    device = density._device if hasattr(density, "_device") else "cpu"
 
     # Setting these is required because _pairplot_scaffold will check if opts['diag'] is
     # `None`. This would break if opts has no key 'diag'. Same for 'upper'.
@@ -653,21 +654,24 @@ def conditional_pairplot(
     opts["lower"] = None
 
     dim, limits, eps_margins = prepare_for_conditional_plot(condition, opts)
-
     diag_func = get_conditional_diag_func(opts, limits, eps_margins, resolution)
 
     def upper_func(row, col, **kwargs):
-        p_image = eval_conditional_density(
-            opts["density"],
-            opts["condition"],
-            limits,
-            row,
-            col,
-            resolution=resolution,
-            eps_margins1=eps_margins[row],
-            eps_margins2=eps_margins[col],
-            warn_about_deprecation=False,
-        ).numpy()
+        p_image = (
+            eval_conditional_density(
+                opts["density"],
+                opts["condition"].to(device),
+                limits.to(device),
+                row,
+                col,
+                resolution=resolution,
+                eps_margins1=eps_margins[row],
+                eps_margins2=eps_margins[col],
+                warn_about_deprecation=False,
+            )
+            .to("cpu")
+            .numpy()
+        )
         h = plt.imshow(
             p_image.T,
             origin="lower",
diff --git a/sbi/analysis/sensitivity_analysis.py b/sbi/analysis/sensitivity_analysis.py
--- a/sbi/analysis/sensitivity_analysis.py
+++ b/sbi/analysis/sensitivity_analysis.py
@@ -27,7 +27,7 @@ def __init__(self, mean: Union[Tensor, float], std: Union[Tensor, float]):
         self.register_buffer("_std", std)
 
     def forward(self, tensor):
-        return tensor * self.std + self.mean
+        return tensor * self._std + self._mean
 
 
 def destandardizing_net(batch_t: Tensor, min_std: float = 1e-7) -> nn.Module:
@@ -110,12 +110,13 @@ def __init__(self, posterior: Any):
             posterior: Posterior distribution obtained with `SNPE`, `SNLE`, or `SNRE`.
                 Needs to have a `.sample()` method. If we want to analyse the
                 sensitivity of the posterior probability, it also must have a
-                `.log_prob()` method.
+                `.potential()` method.
         """
         self._posterior = posterior
         self._regression_net = None
         self._theta = None
         self._emergent_property = None
+        self._device = posterior._device
         self._validation_log_probs = []
 
     def add_property(
@@ -279,7 +280,9 @@ def train(
         )
 
         if self._regression_net is None:
-            self._regression_net = self._build_nn(self._theta[train_indices])
+            self._regression_net = self._build_nn(self._theta[train_indices]).to(
+                self._device
+            )
 
         optimizer = optim.Adam(
             list(self._regression_net.parameters()),
@@ -295,8 +298,8 @@ def train(
             self._regression_net.train()
             for parameters, observations in train_loader:
                 optimizer.zero_grad()
-                outputs = self._regression_net(parameters)
-                loss = criterion(outputs, observations)
+                outputs = self._regression_net(parameters.to(self._device))
+                loss = criterion(outputs, observations.to(self._device))
                 loss.backward()
                 if clip_max_norm is not None:
                     clip_grad_norm_(
@@ -313,8 +316,8 @@ def train(
             val_loss = 0.0
             with torch.no_grad():
                 for parameters, observations in val_loader:
-                    outputs = self._regression_net(parameters)
-                    loss = criterion(outputs, observations)
+                    outputs = self._regression_net(parameters.to(self._device))
+                    loss = criterion(outputs, observations.to(self._device))
                     val_loss += loss.item()
             self._val_log_prob = -val_loss / num_validation_examples
             self._validation_log_probs.append(self._val_log_prob)
@@ -422,7 +425,7 @@ def find_directions(
         thetas.requires_grad = True
 
         if posterior_log_prob_as_property:
-            predictions = self._posterior.log_prob(thetas, track_gradients=True)
+            predictions = self._posterior.potential(thetas, track_gradients=True)
         else:
             predictions = self._regression_net.forward(thetas)
         loss = predictions.mean()
@@ -482,6 +485,7 @@ def project(self, theta: Tensor, num_dimensions: int) -> Tensor:
         Returns:
             Projected parameters of shape `(theta.shape[0], num_dimensions)`.
         """
+        theta = theta.to(self._device)
         if self._gradients_are_normed:
             theta = (theta - self._prior_mean) / self._prior_scale
 
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
@@ -84,11 +84,13 @@ def eval_conditional_density(
         float(limits[dim1, 0] + eps_margins1),
         float(limits[dim1, 1] - eps_margins1),
         resolution,
+        device=condition.device,
     )
     theta_grid_dim2 = torch.linspace(
         float(limits[dim2, 0] + eps_margins2),
         float(limits[dim2, 1] - eps_margins2),
         resolution,
+        device=condition.device,
     )
 
     if dim1 == dim2:
@@ -160,7 +162,7 @@ def conditional_corrcoeff(
             "`from sbi.analysis import conditional_corrcoeff`."
         )
 
-    condition = ensure_theta_batched(condition)
+    device = density._device if hasattr(density, "_device") else "cpu"
 
     if subset is None:
         subset = range(condition.shape[1])
@@ -173,14 +175,14 @@ def conditional_corrcoeff(
                     _compute_corrcoeff(
                         eval_conditional_density(
                             density,
-                            cond,
-                            limits,
+                            cond.to(device),
+                            limits.to(device),
                             dim1=dim1,
                             dim2=dim2,
                             resolution=resolution,
                             warn_about_deprecation=False,
                         ),
-                        limits[[dim1, dim2]],
+                        limits[[dim1, dim2]].to(device),
                     )
                     for dim1 in subset
                     for dim2 in subset
@@ -193,8 +195,10 @@ def conditional_corrcoeff(
 
     # `average_correlations` is still a vector containing the upper triangular entries.
     # Below, assemble them into a matrix:
-    av_correlation_matrix = torch.zeros((len(subset), len(subset)))
-    triu_indices = torch.triu_indices(row=len(subset), col=len(subset), offset=1)
+    av_correlation_matrix = torch.zeros((len(subset), len(subset)), device=device)
+    triu_indices = torch.triu_indices(
+        row=len(subset), col=len(subset), offset=1, device=device
+    )
     av_correlation_matrix[triu_indices[0], triu_indices[1]] = average_correlations
 
     # Make the matrix symmetric by copying upper diagonal to lower diagonal.
@@ -293,7 +297,7 @@ def _expected_value_f_of_x(
     limits = ensure_theta_batched(limits)
 
     x_values_over_which_we_integrate = [
-        torch.linspace(lim[0], lim[1], prob.shape[0])
+        torch.linspace(lim[0], lim[1], prob.shape[0], device=probs.device)
         for lim, prob in zip(torch.flip(limits, [0]), probs)
     ]  # See #403 and #404 for flip().
     grids = list(torch.meshgrid(x_values_over_which_we_integrate))


Test patch:
diff --git a/tests/analysis_test.py b/tests/analysis_test.py
new file mode 100644
--- /dev/null
+++ b/tests/analysis_test.py
@@ -0,0 +1,60 @@
+import torch
+from sbi.inference import SNPE
+from sbi.analysis import ActiveSubspace
+from sbi.utils import BoxUniform
+import pytest
+from sbi.analysis import conditional_pairplot, conditional_corrcoeff
+
+
+@pytest.mark.slow
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ["cpu", "cuda:0"])
+def test_analysis_modules(device: str, set_seed) -> None:
+    """Tests sensitivity analysis and conditional posterior utils on GPU and CPU.
+
+    This test performs only API tests. It does not test the accuracy of the modules.
+
+    Args:
+        device: Which device to run the inference on.
+        set_seed: Fixture for seeding.
+    """
+    num_dim = 3
+    prior = BoxUniform(
+        low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim), device=device
+    )
+
+    def simulator(parameter_set):
+        return 1.0 + parameter_set + torch.randn(parameter_set.shape) * 0.1
+
+    theta = prior.sample((300,)).to("cpu")
+    x = simulator(theta)
+
+    inf = SNPE(prior=prior, device=device)
+    _ = inf.append_simulations(theta, x).train(max_num_epochs=10)
+
+    observation = torch.zeros(num_dim)
+    posterior = inf.build_posterior().set_default_x(observation)
+
+    a = ActiveSubspace(posterior)
+    theta = posterior.sample((200,)).to("cpu")
+    property_ = simulator(theta)[:, :1]
+    _ = a.add_property(theta, property_).train()
+
+    evals, dirs = a.find_directions()
+    assert str(evals.device) == device
+    assert str(dirs.device) == device
+
+    projected = a.project(theta[:10], 1)
+    assert str(projected.device) == device
+
+    # Compute the matrix of correlation coefficients of the slices.
+    cond_coeff_mat = conditional_corrcoeff(
+        density=posterior,
+        condition=posterior.sample((1,)),
+        limits=torch.tensor([[-2.0, 2.0], [-2.0, 2.0], [-2, 2.0]]),
+    )
+    assert str(cond_coeff_mat.device) == device
+
+    _ = conditional_pairplot(
+        posterior, condition=posterior.sample((1,)), limits=[[-2, 2], [-2, 2], [-2, 2]]
+    )
diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py
--- a/tests/inference_on_device_test.py
+++ b/tests/inference_on_device_test.py
@@ -18,6 +18,7 @@
     simulate_for_sbi,
     MCMCPosterior,
     RejectionPosterior,
+    DirectPosterior,
     ratio_estimator_based_potential,
     likelihood_estimator_based_potential,
     posterior_estimator_based_potential,
@@ -32,11 +33,14 @@
     "method, model, mcmc_method",
     [
         (SNPE_C, "mdn", "rejection"),
-        (SNPE_C, "maf", "rejection"),
+        (SNPE_C, "maf", "slice"),
+        (SNPE_C, "maf", "direct"),
         (SNLE, "maf", "slice"),
         (SNLE, "nsf", "slice_np"),
+        (SNLE, "nsf", "rejection"),
         (SNRE_A, "mlp", "slice_np_vectorized"),
         (SNRE_B, "resnet", "nuts"),
+        (SNRE_B, "resnet", "rejection"),
     ],
 )
 @pytest.mark.parametrize("data_device", ("cpu", "cuda:0"))
@@ -136,6 +140,10 @@ def simulator(theta):
                 potential_fn=potential_fn,
                 device=training_device,
             )
+        elif mcmc_method == "direct":
+            posterior = DirectPosterior(
+                posterior_estimator=estimator, prior=prior
+            ).set_default_x(x_o)
         else:
             posterior = MCMCPosterior(
                 potential_fn=potential_fn,