feat: implement posterior.to(device) (#1527)

* add to(device) to BasePotential

* add to(device) to LikelihoodBasedPotential

* add to(device) to PosteriorBasedPotential

* remove Optional for prior

* add to(device) to RatioBasedPotential

* add to(device) to PosteriorScoreBasedPotential

* add to(device) to IIDScoreFunction

* add to(device) to DirectPosterior

* add to(device) to EnsemblePotential and EnsemblePosterior

* add to(device) to ImportanceSamplingPosterior

* add potentials and posterior testing of to(device) method

* add to(device) method to MCMCPosterior

* add to(device) method to RejectionPosterior

* add to(device) to ScorePosterior and corresponding test

* add to(device) to VIPosterior and FIX bug, _prepare_potential is not in this class.

* fix if conditions on attributes

* fix if condition on posteriors that checks if default_x(x_o) has been stored, to keep it

* overcome different types of devices on substraction in transformed_potential

* correct idxs that must be on cpu

* add sample and log_probs test

* put potentials and inits in cpu to calculate map

* remove comments

* revert changes on map

* fix typing for all modified functions and classes

* fix naming from score to vector_field

* fix type checks and score to vector_field naming

* Make sure all tensors are on the same device on VectorFieldPosterior

* include test for all VectorFieldInference samples also on device

* fix moving time tensor to device, and change of inference_method name

* modify assert message

* add device attribute to VectorFieldPosterior in initalization

* fix typo

* Apply suggestions from code review

Co-authored-by: Jan <janfb@users.noreply.github.com>

* fix typing

* fix mnpe device handling

* fix marginal trainer device handling

* fix gradient ascent device handling; refactor tests

* add missinf to(device) to OneDimPriorWrapper

* fix failing test on notImplementedError CustomPriorWrapper

* fix user utils tests

---------

Co-authored-by: Jan <janfb@users.noreply.github.com>
Co-authored-by: Jan Boelts <jan.boelts@mailbox.org>

Author: jorobledo

sbi/inference/posteriors/base_posterior.py

@@ -33,7 +33,7 @@ def __init__(
         self,
         potential_fn: Union[BasePotential, CustomPotential],
         theta_transform: Optional[TorchTransform] = None,
-        device: Optional[str] = None,
+        device: Optional[Union[str, torch.device]] = None,
         x_shape: Optional[torch.Size] = None,
     ):
         """

sbi/inference/posteriors/direct_posterior.py

@@ -46,7 +46,7 @@ def __init__(
         posterior_estimator: ConditionalDensityEstimator,
         prior: Distribution,
         max_sampling_batch_size: int = 10_000,
-        device: Optional[str] = None,
+        device: Optional[Union[str, torch.device]] = None,
         x_shape: Optional[torch.Size] = None,
         enable_transform: bool = True,
     ):
@@ -67,6 +67,8 @@ def __init__(
         # builds it itself. The `potential_fn` and `theta_transform` are used only for
         # obtaining the MAP.
         check_prior(prior)
+        self.enable_transform = enable_transform
+        self.x_shape = x_shape
         potential_fn, theta_transform = posterior_estimator_based_potential(
             posterior_estimator,
             prior,
@@ -81,6 +83,7 @@ def __init__(
             x_shape=x_shape,
         )
 
+        self.device = device
         self.prior = prior
         self.posterior_estimator = posterior_estimator
 
@@ -90,6 +93,45 @@ def __init__(
         self._purpose = """It samples the posterior network and rejects samples that
             lie outside of the prior bounds."""
 
+    def to(self, device: Union[str, torch.device]) -> None:
+        """Move posterior_estimator, prior and x_o to device.
+
+        Changes the device attribute, reinstanciates the
+        posterior, and resets the default x.
+
+        Args:
+            device: device where to move the posterior to.
+        """
+        self.device = device
+        if hasattr(self.prior, "to"):
+            self.prior.to(device)  # type: ignore
+        else:
+            raise ValueError("""Prior has no attribute to(device).""")
+        if hasattr(self.posterior_estimator, "to"):
+            self.posterior_estimator.to(device)
+        else:
+            raise ValueError("""Posterior estimator has no attribute to(device).""")
+
+        potential_fn, theta_transform = posterior_estimator_based_potential(
+            self.posterior_estimator,
+            self.prior,
+            x_o=None,
+            enable_transform=self.enable_transform,
+        )
+        x_o = None
+        if hasattr(self, "_x") and (self._x is not None):
+            x_o = self._x.to(device)
+
+        super().__init__(
+            potential_fn=potential_fn,
+            theta_transform=theta_transform,
+            device=device,
+            x_shape=self.x_shape,
+        )
+        # super().__init__ erases the self._x, so we need to set it again
+        if x_o is not None:
+            self.set_default_x(x_o)
+
     def sample(
         self,
         sample_shape: Shape = torch.Size(),

sbi/inference/posteriors/ensemble_posterior.py

@@ -11,7 +11,7 @@
 from sbi.inference.potentials.base_potential import BasePotential
 from sbi.inference.potentials.posterior_based_potential import PosteriorBasedPotential
 from sbi.sbi_types import Shape, TorchTransform
-from sbi.utils.sbiutils import gradient_ascent
+from sbi.utils.sbiutils import gradient_ascent, mcmc_transform
 from sbi.utils.torchutils import ensure_theta_batched
 from sbi.utils.user_input_checks import process_x
 
@@ -53,13 +53,15 @@ class EnsemblePosterior(NeuralPosterior):
         theta_transform: If passed, this transformation will be applied during the
                 optimization performed when obtaining the map. It does not affect the
                 `.sample()` and `.log_prob()` methods.
+        device: device to host the posterior distribution.
     """
 
     def __init__(
         self,
         posteriors: List,
         weights: Optional[Union[List[float], Tensor]] = None,
         theta_transform: Optional[TorchTransform] = None,
+        device: Optional[Union[str, torch.device]] = None,
     ):
         r"""
         Args:
@@ -77,11 +79,34 @@ def __init__(
         self.theta_transform = theta_transform
         # Take first prior as reference
         self.prior = posteriors[0].potential_fn.prior
+        self.device = device
 
-        device = self.ensure_same_device(posteriors)
+        if self.device is None:
+            self.device = self.ensure_same_device(posteriors)
 
+        self._build_potential_fns()
+
+    def to(self, device: Union[str, torch.device]) -> None:
+        """Moves each posterior to device.
+
+        Prior and weights are also moved to
+        the specified device.
+
+        Args:
+            device: The device to move the ensemble posterior to.
+        """
+        self.device = device
+        self._device = device
+        for i in range(len(self.posteriors)):
+            self.posteriors[i].to(device)
+        self.prior.to(device)
+        self.theta_transform = mcmc_transform(self.prior, device=device)
+        self._weights.to(device)
+        self._build_potential_fns()
+
+    def _build_potential_fns(self):
         potential_fns = []
-        for posterior in posteriors:
+        for posterior in self.posteriors:
             potential = posterior.potential_fn
             potential_fns.append(potential)
             # make sure all prior are the same
@@ -94,8 +119,8 @@ def __init__(
 
         super().__init__(
             potential_fn=potential_fn,
-            theta_transform=theta_transform,
-            device=device,
+            theta_transform=self.theta_transform,
+            device=self.device,
         )
 
     def ensure_same_device(self, posteriors: List) -> str:
@@ -405,7 +430,7 @@ def __init__(
         weights: Tensor,
         prior: Distribution,
         x_o: Optional[Tensor],
-        device: str = "cpu",
+        device: Union[str, torch.device] = "cpu",
     ):
         r"""
         Args:
@@ -419,6 +444,23 @@ def __init__(
         self.potential_fns = potential_fns
         super().__init__(prior, x_o, device)
 
+    def to(self, device: Union[str, torch.device]) -> None:
+        """
+        Moves the ensemble potentials, the prior, the weights and x_o to
+
+        the specified device.
+
+        Args:
+            device: The device to move the ensemble potential to.
+        """
+        self.device = device
+        for i in range(len(self.potential_fns)):
+            self.potential_fns[i].to(device)
+        self._weights = self._weights.to(device)
+        self.prior.to(device)  # type: ignore
+        if self._x_o is not None:
+            self._x_o = self._x_o.to(device)
+
     def allow_iid_x(self) -> bool:
         # in case there is different kinds of posteriors, this will produce an error
         # in `set_x()`

sbi/inference/posteriors/importance_posterior.py

@@ -11,6 +11,7 @@
 from sbi.samplers.importance.importance_sampling import importance_sample
 from sbi.samplers.importance.sir import sampling_importance_resampling
 from sbi.sbi_types import Shape, TorchTransform
+from sbi.utils.sbiutils import mcmc_transform
 from sbi.utils.torchutils import ensure_theta_batched
 
 
@@ -32,7 +33,7 @@ def __init__(
         method: str = "sir",
         oversampling_factor: int = 32,
         max_sampling_batch_size: int = 10_000,
-        device: Optional[str] = None,
+        device: Optional[Union[str, torch.device]] = None,
         x_shape: Optional[torch.Size] = None,
     ):
         """
@@ -65,6 +66,7 @@ def __init__(
         self.proposal = proposal
         self._normalization_constant = None
         self.method = method
+        self.theta_transform = theta_transform
 
         self.oversampling_factor = oversampling_factor
         self.max_sampling_batch_size = max_sampling_batch_size
@@ -74,6 +76,34 @@ def __init__(
             "posterior and can evaluate the _unnormalized_ posterior density with "
             ".log_prob()."
         )
+        self.x_shape = x_shape
+
+    def to(self, device: Union[str, torch.device]) -> None:
+        """
+        Move the potential, the proposal and x_o to a new device.
+
+        It also reinstantiates the posterior with the new device.
+
+        Args:
+            device: Device on which to move the posterior to.
+        """
+        self.device = device
+        self.potential_fn.to(device)  # type: ignore
+        self.proposal.to(device)
+        x_o = None
+        if hasattr(self, "_x") and (self._x is not None):
+            x_o = self._x.to(device)
+
+        self.theta_transform = mcmc_transform(self.proposal, device=device)
+        super().__init__(
+            self.potential_fn,
+            theta_transform=self.theta_transform,
+            device=device,
+            x_shape=self.x_shape,
+        )
+        # super().__init__ erases the self._x, so we need to set it again
+        if x_o is not None:
+            self.set_default_x(x_o)
 
     def log_prob(
         self,
@@ -211,9 +241,9 @@ def sample_batched(
         show_progress_bars: bool = True,
     ) -> Tensor:
         raise NotImplementedError(
-            "Batched sampling is not implemented for ImportanceSamplingPosterior. "
-            "Alternatively you can use `sample` in a loop "
-            "[posterior.sample(theta, x_o) for x_o in x]."
+            "Batched sampling is not implemented for ImportanceSamplingPosterior. \
+           Alternatively you can use `sample` in a loop \
+           [posterior.sample(theta, x_o) for x_o in x]."
         )
 
     def _importance_sample(

sbi/inference/posteriors/mcmc_posterior.py

@@ -59,7 +59,7 @@ def __init__(
         init_strategy_num_candidates: Optional[int] = None,
         num_workers: int = 1,
         mp_context: str = "spawn",
-        device: Optional[str] = None,
+        device: Optional[Union[str, torch.device]] = None,
         x_shape: Optional[torch.Size] = None,
     ):
         """
@@ -136,6 +136,7 @@ def __init__(
         self._posterior_sampler = None
         # Hardcode parameter name to reduce clutter kwargs.
         self.param_name = "theta"
+        self.x_shape = x_shape
 
         if init_strategy_num_candidates is not None:
             warn(
@@ -155,6 +156,34 @@ def __init__(
             "can evaluate the _unnormalized_ posterior density with .log_prob()."
         )
 
+    def to(self, device: Union[str, torch.device]) -> None:
+        """Moves potential_fn, proposal, x_o and theta_transform to the
+
+        specified device. Reinstantiates the posterior and resets the default x_o.
+
+        Args:
+            device: Device to move the posterior to.
+        """
+        self.device = device
+        self.potential_fn.to(device)  # type: ignore
+        self.proposal.to(device)
+        x_o = None
+        if hasattr(self, "_x") and (self._x is not None):
+            x_o = self._x.to(device)
+
+        self.theta_transform = mcmc_transform(self.proposal, device=device)
+
+        super().__init__(
+            self.potential_fn,
+            theta_transform=self.theta_transform,
+            device=device,
+            x_shape=self.x_shape,
+        )
+        # super().__init__ erases the self._x, so we need to set it again
+        if x_o is not None:
+            self.set_default_x(x_o)
+        self.potential_ = self._prepare_potential(self.method)
+
     @property
     def mcmc_method(self) -> str:
         """Returns MCMC method."""
@@ -266,9 +295,9 @@ def sample(
         )
         if init_strategy_num_candidates is not None:
             warn(
-                "Passing `init_strategy_num_candidates` is deprecated as of sbi "
-                "v0.19.0. Instead, use e.g., "
-                f"`init_strategy_parameters={'num_candidate_samples': 1000}`",
+                f"Passing `init_strategy_num_candidates` is deprecated as of sbi \
+                v0.19.0. Instead, use e.g., \
+                `init_strategy_parameters={'num_candidate_samples': 1000}`",
                 stacklevel=2,
             )
             self.init_strategy_parameters["num_candidate_samples"] = (