Prior to(device) (#1505)

* add to(device) method for BoxUniform and test

* add to(device) to PytorchReturnTypeWrapper

* add process device to tests

* add to(device) for MultipleIndependent

* fix names, docstring, and not implemented to(device) on CurstomPriorWrapper

* add type hints and process_prior for BoxUniform

* add torch.device to process_device

* add type hints for functions in user_input_checks_utils.py

* add docstring for prior '.to()' methods

* type hint added in BoxUniform

* add test for MultivariateNormal, Binomial, and Normal. Also fix format.

* add assertion messages in prior tests

* add torch.device in dech_device

* fix type of dist in get_distributions_parameters

* add docstrings for prior_device_tests

* change return type of process_device to also return torch.device

* fix torch.device type when processing

* add test for torch.device as input of process_device

---------

Co-authored-by: vepe99 <viterbogiuseppe99@gmail.com>

Author: jorobledo

sbi/utils/torchutils.py

@@ -15,7 +15,7 @@
 from sbi.utils.typechecks import is_nonnegative_int, is_positive_int
 
 
-def process_device(device: str) -> str:
+def process_device(device: Union[str, torch.device]) -> str:
     """Set and return the default device to cpu or gpu (cuda, mps).
 
     Args:
@@ -52,6 +52,8 @@ def process_device(device: str) -> str:
         # Else, check whether the custom device is valid.
         else:
             check_device(device)
+            if isinstance(device, torch.device):
+                device = device.type
 
         return device
 
@@ -61,7 +63,7 @@ def gpu_available() -> bool:
     return torch.cuda.is_available() or torch.backends.mps.is_available()
 
 
-def check_device(device: str) -> None:
+def check_device(device: Union[str, torch.device]) -> None:
     """Check whether the device is valid.
 
     Args:
@@ -274,8 +276,8 @@ def gaussian_kde_log_eval(samples, query):
 class BoxUniform(Independent):
     def __init__(
         self,
-        low: Tensor,
-        high: Tensor,
+        low: Union[Tensor, Array],
+        high: Union[Tensor, Array],
         reinterpreted_batch_ndims: int = 1,
         device: Optional[str] = None,
     ):
@@ -313,6 +315,11 @@ def __init__(
         # Device handling
         device = low.device.type if device is None else device
         device = process_device(device)
+        self.device = device
+        self.reinterpreted_batch_ndims = reinterpreted_batch_ndims
+
+        self.low = torch.as_tensor(low, dtype=torch.float32, device=device)
+        self.high = torch.as_tensor(high, dtype=torch.float32, device=device)
 
         super().__init__(
             Uniform(
@@ -327,6 +334,30 @@ def __init__(
             reinterpreted_batch_ndims,
         )
 
+    def to(self, device: Union[str, torch.device]) -> None:
+        """
+        Moves the distribution to the specified device **in place**.
+
+        Args:
+            device: Target device (e.g., "cpu", "cuda", "mps").
+        """
+        # Update the device attribute
+        self.device = device
+        device = process_device(device)
+
+        # Move tensors to the new device
+        self.low = self.low.to(device=device)
+        self.high = self.high.to(device=device)
+
+        super().__init__(
+            Uniform(
+                low=self.low,
+                high=self.high,
+                validate_args=False,
+            ),
+            self.reinterpreted_batch_ndims,
+        )
+
 
 def ensure_theta_batched(theta: Tensor) -> Tensor:
     r"""

sbi/utils/user_input_checks_utils.py

@@ -2,13 +2,35 @@
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
 import warnings
-from typing import Dict, Optional, Sequence
+from typing import Dict, Optional, Sequence, Union
 
 import torch
 from torch import Tensor, float32
 from torch.distributions import Distribution, constraints
 
 
+def get_distribution_parameters(
+    dist: torch.distributions.Distribution, device: Union[str, torch.device]
+) -> Dict:
+    """Used to get the tensors of the parameters in torch distributions.
+
+    Returns the tensors relocated to device.
+    """
+    params = {param: getattr(dist, param).to(device) for param in dist.arg_constraints}
+    # torch.distributions.MultivariateNormal calculates precision
+    # matrix from covariance, and stores it in the arg_constraints.
+    # When reinstantiating, we must provide only one of them.
+    if isinstance(dist, torch.distributions.MultivariateNormal):
+        params["precision_matrix"] = None
+        params["scale_tril"] = None
+    # torch.distributions.MultivariateNormal calculates logits
+    # from probabilities, and stores it in the arg_constraints.
+    # When reinstantiating, we must provide only one of them.
+    elif isinstance(dist, torch.distributions.Binomial):
+        params["logits"] = None
+    return params
+
+
 class CustomPriorWrapper(Distribution):
     def __init__(
         self,
@@ -83,6 +105,18 @@ def _set_mean_and_variance(self):
                 stacklevel=2,
             )
 
+    def to(self, device: Union[str, torch.device]) -> None:
+        """
+        Move the distribution to the specified device. Not implemented for this class.
+
+        Raises:
+            NotImplementedError.
+        """
+        raise NotImplementedError(
+            "This class is not supported on the GPU. Use on cpu or use \
+            any of `PytorchReturnTypeWrapper`, `BoxUniform`, or `MultipleIndependent`."
+        )
+
     @property
     def mean(self):
         return torch.as_tensor(
@@ -118,6 +152,7 @@ def __init__(
         )
 
         self.prior = prior
+        self.device = None
         self.return_type = return_type
 
     def log_prob(self, value) -> Tensor:
@@ -150,6 +185,20 @@ def variance(self):
     def support(self):
         return self.prior.support
 
+    def to(self, device: Union[str, torch.device]) -> None:
+        """
+        Move the distribution to the specified device.
+
+        Moves the distribution parameters to the specific device
+        and updates the device attribute.
+
+        Args:
+            device: device to move the distribution to.
+        """
+        params = get_distribution_parameters(self.prior, device)
+        self.prior = type(self.prior)(**params)
+        self.device = device
+
 
 class MultipleIndependent(Distribution):
     """Wrap a sequence of PyTorch distributions into a joint PyTorch distribution.
@@ -181,6 +230,7 @@ def __init__(
             [d.set_default_validate_args(validate_args) for d in dists]
 
         self.dists = dists
+        self.device = None
         # numel() instead of event_shape because for all dists both is possible,
         # event_shape=[1] or batch_shape=[1]
         self.dims_per_dist = [d.sample().numel() for d in self.dists]
@@ -319,6 +369,24 @@ def support(self):
             reinterpreted_batch_ndims=1,
         )
 
+    def to(self, device: Union[str, torch.device]) -> None:
+        """
+        Move the distribution to the specified device.
+        If the distribution has the `to` method, it is used. Otherwise, the
+        parameters of the distribution are moved to the specified device.
+
+        Args:
+            device: device to move the distribution to.
+        """
+        for i in range(len(self.dists)):
+            # ignoring because it is related to torch and not sbi
+            if hasattr(self.dists[i], "to"):
+                self.dists[i].to(device)  # type: ignore
+            else:
+                params = get_distribution_parameters(self.dists[i], device)
+                self.dists[i] = type(self.dists[i])(**params)  # type: ignore
+        self.device = device
+
 
 def build_support(
     lower_bound: Optional[Tensor] = None, upper_bound: Optional[Tensor] = None

tests/inference_on_device_test.py

@@ -444,12 +444,15 @@ def allow_iid_x(self) -> bool:
         ("gpu", None),
         ("cpu", "cpu"),
         ("gpu", "gpu"),
+        (torch.device("cpu"), torch.device("cpu")),
         pytest.param("gpu", "cpu", marks=pytest.mark.xfail),
         pytest.param("cpu", "gpu", marks=pytest.mark.xfail),
     ],
 )
 def test_boxuniform_device_handling(arg_device, device):
-    """Test mismatch between device passed via low / high and device kwarg."""
+    """Test mismatch between device passed via low / high and device kwarg.
+
+    Also tests torch.device as argument of process_device."""
 
     arg_device = process_device(arg_device)
     device = process_device(device)

tests/prior_device_tests.py

@@ -0,0 +1,106 @@
+import pytest
+import torch
+from torch.distributions import Beta, Binomial, Gamma, MultivariateNormal, Normal
+
+from sbi.utils.torchutils import BoxUniform, process_device
+from sbi.utils.user_input_checks_utils import (
+    MultipleIndependent,
+    PytorchReturnTypeWrapper,
+)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ["cpu", "gpu"])
+def test_BoxUniform(device: str):
+    """Test moving BoxUniform prior between devices."""
+    device = process_device(device)
+    low = torch.tensor([0.0])
+    high = torch.tensor([1.0])
+    prior = BoxUniform(low, high)
+    sample = prior.sample((1,))
+    assert prior.device == "cpu", "Prior is not initially in cpu."
+    assert sample.device.type == "cpu", "sample is not initially in cpu."
+    log_probs = prior.log_prob(sample)
+    assert log_probs.device.type == "cpu", "Log probs are not initially in cpu."
+
+    prior.to(device)
+    assert prior.device == device, f"Prior was not moved to {device}."
+    assert prior.low.device.type == device.strip(":0"), (
+        f"BoxUniform low tensor is not in {device}."
+    )
+    assert prior.high.device.type == device.strip(":0"), (
+        f"BoxUniform high tensor is not in {device}."
+    )
+
+    sample_device = prior.sample((100,))
+    assert sample_device.device.type == device.strip(":0"), (
+        f"sample tensor is not in {device}."
+    )
+    log_probs = prior.log_prob(sample_device)
+    assert log_probs.device.type == device.strip(":0"), (
+        f"log_prob tensor is not in {device}."
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ["cpu", "gpu"])
+@pytest.mark.parametrize(
+    "prior",
+    [
+        Normal(loc=0.0, scale=1.0),
+        Binomial(total_count=10, probs=torch.tensor([0.5])),
+        MultivariateNormal(torch.tensor([0.1, 0.0]), covariance_matrix=torch.eye(2)),
+    ],
+)
+def test_PytorchReturnTypeWrapper(device: str, prior: torch.distributions):
+    """Test moving PytorchReturnTypeWrapper objects between devices.
+
+    Asserts that samples, prior, and log_probs are in device.
+    """
+    device = process_device(device)
+    prior = PytorchReturnTypeWrapper(prior)
+
+    prior.to(device)
+    assert prior.device == device, f"Prior was not correctly moved to {device}."
+
+    sample_device = prior.sample((100,))
+    assert sample_device.device.type == device.strip(":0"), (
+        f"sample was not correctly moved to {device}."
+    )
+    log_probs = prior.log_prob(sample_device)
+    assert log_probs.device.type == device.strip(":0"), (
+        f"log_prob was not correctly moved to {device}."
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ["cpu", "gpu"])
+def test_MultipleIndependent(device: str):
+    """Test moving MultipleIndependent objects between devices.
+
+    Asserts that samples, prior, and log_probs are in device.
+    Uses Gamma, Beta, Normal and Binomial, from
+    torch.distributions and BoxUniform form sbi.
+    """
+    device = process_device(device)
+    dists = [
+        Gamma(torch.tensor([1.0]), torch.tensor([0.5])),
+        Beta(torch.tensor([2.0]), torch.tensor([2.0])),
+        BoxUniform(torch.zeros(1), torch.ones(1)),
+        Normal(torch.tensor([0.0]), torch.tensor([0.5])),
+        Binomial(torch.tensor([10]), torch.tensor([0.5])),
+    ]
+
+    prior = MultipleIndependent(dists)
+
+    prior.to(device)
+    assert prior.device == device, f"Prior was not correctly moved to {device}."
+
+    sample_device = prior.sample((100,))
+    assert sample_device.device.type == device.strip(":0"), (
+        f"sample was not correctly moved to {device}."
+    )
+    log_probs = prior.log_prob(sample_device)
+    assert log_probs.device.type == device.strip(":0"), (
+        f"log_prob was not correctly moved to {device}."
+    )