posterior.to(device)

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 @@
         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 @@
         # 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 @@
             x_shape=x_shape,
         )
 
+        self.device = device
         self.prior = prior
         self.posterior_estimator = posterior_estimator
 
@@ -90,6 +93,45 @@
         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__ erase 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 @@
         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 @@
         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 @@
 
         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 @@
         weights: Tensor,
         prior: Distribution,
         x_o: Optional[Tensor],
-        device: str = "cpu",
+        device: Union[str, torch.device] = "cpu",
     ):
         r"""
         Args:
@@ -419,6 +444,23 @@
         self.potential_fns = potential_fns
         super().__init__(prior, x_o, device)
 
+    def to(self, device):
+        """
+        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 @@
         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 @@
         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 @@
             "posterior and can evaluate the _unnormalized_ posterior density with "
             ".log_prob()."
         )
+        self.x_shape = x_shape
+
+    def to(self, device):
+        """
+        Move the potential, the proposal and x_o to a new device.
+
+        It also reinstansiates 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__ erase 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 @@
         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 @@
         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 @@
         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 @@
             "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
+
+        speficied device. Reinstanciates 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__ erase 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 @@
         )
         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"] = (

sbi/inference/posteriors/rejection_posterior.py

@@ -12,6 +12,7 @@
 from sbi.inference.potentials.base_potential import BasePotential
 from sbi.samplers.rejection.rejection import rejection_sample
 from sbi.sbi_types import Shape, TorchTransform
+from sbi.utils import mcmc_transform
 from sbi.utils.torchutils import ensure_theta_batched
 
 
@@ -24,14 +25,14 @@
 
     def __init__(
         self,
-        potential_fn: Union[Callable, BasePotential],
+        potential_fn: Union[Callable, BasePotential],  # type: ignore
         proposal: Any,
         theta_transform: Optional[TorchTransform] = None,
         max_sampling_batch_size: int = 10_000,
         num_samples_to_find_max: int = 10_000,
         num_iter_to_find_max: int = 100,
         m: float = 1.2,
-        device: Optional[str] = None,
+        device: Optional[Union[str, torch.device]] = None,
         x_shape: Optional[torch.Size] = None,
     ):
         """
@@ -64,12 +65,40 @@
         self.num_samples_to_find_max = num_samples_to_find_max
         self.num_iter_to_find_max = num_iter_to_find_max
         self.m = m
+        self.x_shape = x_shape
 
         self._purpose = (
             "It provides rejection sampling to .sample() from the posterior and "
             "can evaluate the _unnormalized_ posterior density with .log_prob()."
         )
 
+    def to(self, device: Union[str, torch.device]) -> None:
+        """
+        Move potential fucntion, proposal and x_o to the device.
+
+        This method reinstanciate the posterior and resets the default x_o
+
+        Args:
+            device: The 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__ erase 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, theta: Tensor, x: Optional[Tensor] = None, track_gradients: bool = False
     ) -> Tensor: